trajectory_sampling_gui.py

Go to the documentation of this file.

#!/usr/bin/env python

PKG = 'trajectory_sampling_gui'
import roslib; roslib.load_manifest(PKG)

import rospy
import rosbag

import os
import optparse
import sys
import time
import thread
import copy

from sensor_msgs.msg import JointState
from dmp_motion_learner.srv import LearnJointSpaceDMP
from dmp_motion_controller.srv import AddToExecuteDMPQueue, AddToDualArmExecuteDMPQueue
from pr2_mechanism_msgs.srv import SwitchController

import wxversion
WXVER = '2.8'
if wxversion.checkInstalled(WXVER):
    wxversion.select(WXVER)
else:
    print >> sys.stderr, 'This application requires wxPython version %s' % WXVER
    sys.exit(1)
import wx

class DMP:
    LEFT = 0
    RIGHT = 1
    BOTH = 2

class TrajectorySamplingGuiApp(wx.App):
    def __init__(self, options):
        self.options = options
        self.bag = None
        self.sub = None
        wx.App.__init__(self)
    
    def OnInit(self):
        frame = TrajectorySamplingFrame(None, self.options)
        frame.Show()
        self.SetTopWindow(frame)
        return True

class TrajectorySamplingFrame(wx.Frame):
    def __init__(self, parent, options, id=wx.ID_ANY, title='TrajectorySampling', pos=wx.DefaultPosition, size=(320, 330), style=wx.DEFAULT_FRAME_STYLE):
        wx.Frame.__init__(self, parent, id, title, pos, size, style)

        self.options = options

        self.font = wx.Font(9, wx.FONTFAMILY_MODERN, wx.FONTSTYLE_NORMAL, wx.FONTWEIGHT_NORMAL)

        self.colors = ['red', 'blue', 'purple', 'cyan', 'brown']
        self.robots = ['sim', 'pr2']
        self.robot = 'sim'
        
        self.recorders = ['joint_state', 'external']
        self.recorder = 'joint_state'

        self.what = 'sample_trajectory_to_bag_file'

        self.filename = None
        self.recording = False
        self.first_time = False
        self.record_pid = None
        self.mark_sub = None
        
        self.init_from_rospy()
        
        # setup services
        # enable rl loop if requested
        if options.rl:
            from learner_node.srv import start
            from learner_node.srv import stop
            from learner_node.srv import setStartState
            from rtrl.srv import setGoal
            self.learner_start_call = rospy.ServiceProxy('/learner_node/start', start)
            self.learner_stop_call = rospy.ServiceProxy('/learner_node/stop', stop)
            self.learner_setstart_call = rospy.ServiceProxy('/learner_node/set_start_state', setStartState)
            self.dmp_setgoal_call = rospy.ServiceProxy('/rtrl_controller/dmp/set_goal', setGoal)
            
        self.switch_controller = rospy.ServiceProxy('pr2_controller_manager/switch_controller', SwitchController)
        self.dmp_call = rospy.ServiceProxy('/dmp_motion_learner/learn_joint_space_dmp_from_bag_file', LearnJointSpaceDMP)
        self.execute_dmp_left  = rospy.ServiceProxy('/l_arm_dmp_joint_position_controller/add_to_execute_dmp_queue', AddToExecuteDMPQueue)
        self.execute_dmp_right = rospy.ServiceProxy('/r_arm_dmp_joint_position_controller/add_to_execute_dmp_queue', AddToExecuteDMPQueue)
        self.execute_dmp_both  = rospy.ServiceProxy('/dual_arm_dmp_joint_position_controller/add_to_execute_dmp_queue', AddToDualArmExecuteDMPQueue)
        self.dmp_execute_call  = 1

        self.build_menu()
        self.build_main_gui()
        
        # --> try to list all bag files
        self.fill_bag_chooser()
    
    ## initialization
    
    def init_from_rospy(self):
        
        # default type --> right arm
        self.dmp_type = DMP.RIGHT
        
        self.dmp_joints_r = ['r_shoulder_pan_joint', 'r_shoulder_lift_joint', 'r_upper_arm_roll_joint', 'r_elbow_flex_joint', 'r_forearm_roll_joint', 'r_wrist_flex_joint', 'r_wrist_roll_joint']
        self.dmp_joints_l = ['l_shoulder_pan_joint', 'l_shoulder_lift_joint', 'l_upper_arm_roll_joint', 'l_elbow_flex_joint', 'l_forearm_roll_joint', 'l_wrist_flex_joint', 'l_wrist_roll_joint']
        self.joints_r_arm = ['r_upper_arm_roll_joint', 'r_shoulder_pan_joint','r_shoulder_lift_joint','r_forearm_roll_joint','r_elbow_flex_joint','r_wrist_flex_joint','r_wrist_roll_joint']
        self.joints_l_arm = ['l_upper_arm_roll_joint', 'l_shoulder_pan_joint','l_shoulder_lift_joint','l_forearm_roll_joint','l_elbow_flex_joint','l_wrist_flex_joint','l_wrist_roll_joint']
        
        self.start_state_planner_r = []
        self.start_state_planner_l = []
        
        if not self.options.offline:
            if not rospy.has_param('/trajectory_sampling/data_directory_name'):
                rospy.logerr("/trajectory_sampling/data_directory_name not set!")
                self.bag_dir = os.path.expanduser("~") +  "/bag_dir" # sensible default ;)
            else:
                self.bag_dir = rospy.get_param('/trajectory_sampling/data_directory_name')
                
            if not rospy.has_param('/dmp_execution/dmp_dir'):
                rospy.logerr("/dmp_execution/dmp_dir not set!")
                self.scp_bag_dir = os.path.expanduser("~") +  "/bag_dir"
            else:
                self.scp_bag_dir = rospy.get_param('/dmp_execution/dmp_dir')
            
            if rospy.has_param('/learner_node/start_state_names'):
                rospy.logerr("learner_node not set!")
                self.tracking_joints = set(rospy.get_param('/learner_node/start_state_names'))
            else:
                self.tracking_joints = set(self.joints_r_arm)
        else:
            self.bag_dir = os.path.expanduser("~") + "/bag_dir"
            self.scp_bag_dir = os.path.expanduser("~") + "/bag_dir"
            self.tracking_joints = set(self.joints_r_arm)
    
    def build_menu(self):
        self.menu_bar = wx.MenuBar()

        self.file_menu = wx.Menu()
        self.file_exit_item = wx.MenuItem(self.file_menu, -1, 'E&xit', 'Exit the program')
        self.Bind(wx.EVT_MENU, self.on_exit, self.file_exit_item)
        self.file_menu.AppendItem(self.file_exit_item)
        self.menu_bar.Append(self.file_menu, '&File')
        
        self.dmp_menu = wx.Menu()
        self.type_r_arm_item = wx.MenuItem(self.dmp_menu, -1, '&right arm', 'Use right arm for dmp')
        self.Bind(wx.EVT_MENU, self.on_choose_right, self.type_r_arm_item)
        self.dmp_menu.AppendItem(self.type_r_arm_item)
        self.type_l_arm_item = wx.MenuItem(self.dmp_menu, -1, '&light arm', 'Use left arm for dmp')
        self.Bind(wx.EVT_MENU, self.on_choose_left, self.type_l_arm_item)
        self.dmp_menu.AppendItem(self.type_l_arm_item)
        self.type_b_arm_item = wx.MenuItem(self.dmp_menu, -1, '&both arms', 'Use both arms for dmp')
        self.Bind(wx.EVT_MENU, self.on_choose_both, self.type_b_arm_item)
        self.dmp_menu.AppendItem(self.type_b_arm_item)
        self.menu_bar.Append(self.dmp_menu, '&Type')
        
        self.SetMenuBar(self.menu_bar)

    def build_main_gui(self):
        self.output_text = wx.TextCtrl  (self, -1, '', pos=(1, 100), size=(318, 130), style=wx.TE_MULTILINE)
        self.record_bag_text = wx.TextCtrl  (self, -1, 'record.bag', pos=(5, 5), size=(180, 29))
        self.dmp_bag_text = wx.TextCtrl  (self, -1, 'dmp.bag', pos=(5, 30), size=(180, 29))
        self.bag_chooser = wx.ComboBox  (self, -1, 'Recorded Bag Files', pos=(5, 59), size=(180, 29), style=wx.CB_READONLY)
        self.record_button = wx.Button    (self, -1, 'Start Record', pos=(200, 5))
        self.learn_button = wx.Button    (self, -1, 'Learn DMP',             pos=(200,  35))
        self.distribute_button = wx.Button    (self, -1, 'Distribute DMP',             pos=(200,  65))
        if options.rl:
            # when rl is activated show more options to control the learning loop
            self.start_button = wx.Button    (self, -1, 'Start Execution',             pos=(30,  240))
            self.stop_button = wx.Button    (self, -1, 'Stop Execution',             pos=(170,  240))
            self.mark_start_button = wx.Button    (self, -1, 'Mark Start',             pos=(30,  270))
            self.mark_end_button = wx.Button    (self, -1, 'Mark Goal',             pos=(170,  270))
            self.Bind(wx.EVT_BUTTON, lambda e: self.start_learner(), self.start_button)
            self.Bind(wx.EVT_BUTTON, lambda e: self.stop_learner(), self.stop_button)
            self.Bind(wx.EVT_BUTTON, lambda e: self.mark_start(), self.mark_start_button)
            self.Bind(wx.EVT_BUTTON, lambda e: self.mark_end(), self.mark_end_button)
        else:
            # otherwise allow for simple playback of the recorded dmp
            self.start_button = wx.Button    (self, -1, 'Execute DMP',             pos=(105,  240))
            self.Bind(wx.EVT_BUTTON, lambda e: self.execute_dmp(), self.start_button)
        
        self.output_text.SetFont(self.font)
        #self.color_combo = wx.ComboBox  (self, -1, '', pos=(240, 80), size=(90, 29), style=wx.CB_READONLY)
        
        self.record_button.SetBackgroundColour(wx.GREEN)

        #self.Bind(wx.EVT_SIZE, self.on_size)

        #self.output_text.Bind(wx.EVT_TEXT, self.on_text)
        
        self.Bind(wx.EVT_BUTTON, lambda e: self.toggle_recording(), self.record_button)
        self.Bind(wx.EVT_BUTTON, lambda e: self.learn(), self.learn_button)
        self.Bind(wx.EVT_BUTTON, lambda e: self.distribute(), self.distribute_button)
    
    def fill_bag_chooser(self):
        bags = os.listdir(self.bag_dir)
        self.bag_chooser.SetItems(sorted(bags))
    
    ## helpers
    def append_output(self, output):
        self.output_text.AppendText(output + '\n')
    
    def on_text(self, event):
        # Scroll to end
        self.output_text.ScrollPages(1)
    
    ## callbacks
    
    def on_exit(self, event):
        self.Close()
        
    def on_choose_right(self, event):
        self.dmp_type = DMP.RIGHT
    
    def on_choose_left(self, event):
        self.dmp_type = DMP.LEFT
    
    def on_choose_both(self, event):
        self.dmp_type = DMP.BOTH
    
    ## Recording

    def toggle_recording(self):
        if self.recording:
            self.stop_recording()
        else:
            self.start_recording()
    
    def learn(self):
        bag_file = str(self.bag_chooser.GetValue())
        self.learn_dmp(bag_file)
    
    def learn_dmp(self, filename):
        # --> default joints to right 
        dmp_joints = self.dmp_joints_r[:]
        # and adapt if necessary
        if self.dmp_type == DMP.LEFT:
            dmp_joints = self.dmp_joints_l[:]
        elif self.dmp_type == DMP.BOTH:
            dmp_joints.extend(self.dmp_joints_l)

        resp = self.dmp_call(joint_names = dmp_joints, bag_file_name = self.bag_dir + '/' + filename, data_directory_name = self.bag_dir, dmp_id = 1)

        self.append_output(str(resp))

        # store the dmp to a bag file
        bag = rosbag.Bag(self.bag_dir + '/' + str(self.dmp_bag_text.GetValue()), 'w')
        bag.write('dynamic_movement_primitive', resp.dmp)
        bag.close()
    
    def execute_dmp(self):
        exec_dmp = None
        bag = rosbag.Bag(self.bag_dir + '/' + str(self.dmp_bag_text.GetValue()))
        for topic, msg, t in bag.read_messages(topics=['dynamic_movement_primitive']):
            if topic == 'dynamic_movement_primitive': # safety check
                exec_dmp = msg
        bag.close()
        if exec_dmp == None:
            self.append_output('Error could not read dmp from bag!')
            return
        exec_dmp.is_setup = True
        exec_dmp.is_start_set = False
        if self.dmp_type == DMP.RIGHT:
            self.switch_controller.call(start_controllers = ['r_arm_dmp_joint_position_controller'], stop_controllers = ['r_arm_controller'], strictness = 2)
            self.execute_dmp_right.call(dmps = [exec_dmp], execution_durations = [exec_dmp.initial_duration], types = [0])
        elif self.dmp_type == DMP.LEFT:
            self.switch_controller.call(start_controllers = ['l_arm_dmp_joint_position_controller'], stop_controllers = ['l_arm_controller'], strictness = 2)
            self.execute_dmp_left.call(dmps = [exec_dmp], execution_durations = [exec_dmp.initial_duration], types = [0])
        elif self.dmp_type == DMP.BOTH:
            self.switch_controller.call(start_controllers = ['dual_arm_dmp_joint_position_controller', 'r_arm_dmp_joint_position_controller', 'l_arm_dmp_joint_position_controller', ], stop_controllers = ['r_arm_controller', 'l_arm_controller'], strictness = 2)
            right_exec = copy.deepcopy(exec_dmp)
            right_exec.transformation_systems = right_exec.transformation_systems[0:7]
            left_exec = copy.deepcopy(exec_dmp)
            left_exec.transformation_systems = left_exec.transformation_systems[7:]
            self.execute_dmp_both.call(left_arm_dmps = [left_exec], right_arm_dmps = [right_exec], execution_durations = [exec_dmp.initial_duration], left_arm_types = [0], right_arm_types = [0])
    
    def distribute(self):
                def copy_dmp(source, target):
                        print("DISTRIBUTING dmp")
                        os.system('scp %s marvin:%s' % (source, target))
                        #os.system('scp tt localhost:~/')
                        print("DONE")
                        
                self.append_output('~~~ DISTRIBUTING dmp ~~~\n--> see command line for instructions')
                source = self.bag_dir + '/' + str(self.dmp_bag_text.GetValue())
                target = self.scp_bag_dir + '/' + self.dmp_bag_text.GetValue()
                thread.start_new_thread(copy_dmp ,(source, target))
                
                if not self.options.offline:
                        rospy.set_param('/dmp_execution/dmp_file_name', target)
                #os.system('echo "scp %s marvin:%s/%s"' % (source, self.scp_bag_dir, self.dmp_bag_text.GetValue()))
    
    def start_learner(self):
                call_string = '/learner_node/start'
                self.append_output('calling %s' % (call_string))
                resp = self.learner_start_call(rospy.Time.now())
                self.append_output('Starting learner resp: %s' % (str(resp)))
    
    def stop_learner(self):
                call_string = '/learner_node/stop'
                self.append_output('calling %s' % (call_string))
                resp = self.learner_stop_call(rospy.Time.now())
                self.append_output('Stopping learner resp: %s' % (str(resp)))
        
    def mark_start(self):
        if self.mark_sub:
            self.mark_sub.unregister()
        self.append_output('Registering for sampling start state')
        self.mark_sub = rospy.Subscriber('/joint_states', JointState, self.mark_start_callback)
        
    def mark_start_callback(self, msg):
        idx = 0
        start_state = []
        for name in msg.name:
            if name in self.tracking_joints:
                # pos
                start_state.append(msg.position[idx])
                # vel
                start_state.append(0.0)
                # acc
                start_state.append(0.0)
            idx += 1
        # unregister --> we only need to sample one state here
        self.mark_sub.unregister()
        rospy.set_param('/learner_node/start_state', start_state)
        self.append_output('Done sampling start state --> call learner')
        resp = self.learner_setstart_call(rospy.Time.now(), start_state);
        self.append_output('learner response: %s' % (str(resp)))
        
    def mark_end(self):
        if self.mark_sub:
            self.mark_sub.unregister()
        self.append_output('Registering for sampling goal state')
        self.mark_sub = rospy.Subscriber('/joint_states', JointState, self.mark_end_callback)
    
    def mark_end_callback(self, msg):
        idx = 0
        end_state = []
        state_dict = dict()
        for name in msg.name:
            state_dict[name] = msg.position[idx]
            idx += 1
        for name in self.dmp_joints:  
            if name in state_dict:
                # pos
                end_state.append(state_dict[name])
        # unregister --> we only need to sample one state here
        self.mark_sub.unregister()
        rospy.set_param('/learner_node/terminal_pose', end_state)
        self.append_output('Done sampling goal state --> call dmp')
        resp = self.dmp_setgoal_call(rospy.Time.now(), end_state);
        self.append_output('dmp response: %s' % (str(resp)))
    
    def start_recording(self):
        if self.recording:
            return
        
        self.first_time = True
        self.record_button.SetLabel('Stop Record')
        self.record_button.SetBackgroundColour(wx.RED)
        
        self.filename = self.record_bag_text.GetValue()
                
        self.recording = True
        
        self.bag = rosbag.Bag(self.bag_dir + '/' + self.filename, 'w')
        self.sub = rospy.Subscriber('/joint_states', JointState, self.record_callback)

    def record_callback(self, msg):
        if self.first_time:
            self.first_time = False
            self.start_state_planner_r = []
            self.start_state_planner_l = []
            idx = 0
            for name in msg.name:
                if name in self.joints_r_arm:
                    self.start_state_planner_r.append(msg.position[idx])
                if name in self.joints_l_arm:
                    self.start_state_planner_l.append(msg.position[idx])
                idx += 1
            
        self.bag.write('/joint_states', msg)

    def stop_recording(self):
        if not self.recording:
            return
        
        self.sub.unregister()
        
        # wait for bag writing to finish
        time.sleep(0.1)
        self.bag.close()
        
        self.record_button.SetLabel('Start Record')
        self.record_button.SetBackgroundColour(wx.GREEN)

        # Wait before copying the file
        time.sleep(0.1)

        # re check for all bag files
        self.fill_bag_chooser()

        self.recording = False


if __name__ == '__main__':
    parser = optparse.OptionParser()
    parser.add_option('-o', '--offline', action='store_true', default=False, help='run offline - don\'t connect to core')   
    parser.add_option('-r', '--rl', action='store_true', default=False, help='enable reinforcement learning loop --> rtrl package required')
    options, args = parser.parse_args(sys.argv[1:])

    if not options.offline:
        rospy.init_node('trajectory_sampling_gui', anonymous=True)

    app = TrajectorySamplingGuiApp(options)
    app.MainLoop()
 
    rospy.signal_shutdown('GUI shutdown')