pomdp_execution_for_models.py

Go to the documentation of this file.

#!/usr/bin/env python

import roslib; roslib.load_manifest('estirabot_apps')
import rospy
import smach
import smach_ros
import tf
import random

from smach         import StateMachine, CBState, Iterator
from smach_ros     import ServiceState
from actionlib     import *
from actionlib.msg import *
from std_msgs.msg  import String, Int8, Int32
from object_manipulation_msgs.msg import PickupAction,PickupGoal,PlaceAction,PlaceGoal
from sensor_msgs.msg import PointCloud2, JointState
from iri_deformable_analysis.srv import DoDeformableAnalysis, DoDeformableAnalysisRequest
from estirabot_msgs.msg import PomdpGraspConfig
from estirabot_msgs.srv import obs2action, obs2actionRequest
from geometry_msgs.msg import Transform, PoseStamped, Pose
from pprint import pprint

from iri_common_smach.st_get_pcl          import GetPCL
from iri_common_smach.utils_msg           import build_pose_stamped_msg
from iri_common_smach.get_keystroke       import GetKeyStroke
from iri_common_smach.homogeneous_product import homogeneous_product_pose_transform
from iri_common_smach.st_topic_publisher  import TopicPublisher
from iri_wam_smach.sm_wam_move_from_pose  import SM_WAM_MoveFromPose

from estirabot_apps_base.utils_pomdp            import PomdpConfigFactory, TranslatePomdp
from estirabot_apps_base.sm_move                import SM_ESTIRABOT_GoHome
from estirabot_apps_base.sm_deformable_analysis import SM_ESTIRABOT_DeformableAnalysis
from estirabot_apps_base.sm_generate_grasp      import SM_ESTIRABOT_GenerateGraspMsg
from TranslatePomdp2piles import TranslatePomdp 

class ExperimentData():
    def __init__(self):
        ########## Pol rocks and stones ##########################
        # how to create an index that is consistent between executions?
        self.iteration_number = 0 

        self.algorithms = [    1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1,
        1, 1, 1] 
##self.algorithms = [   1,   2,   3,   4,   5,   6, 
#                             11,  12,  13,  14,  15,  16,
#                            101, 102, 103, 104, 105, 106,
#                            111, 112, 113, 114, 115, 116]

        self.algorithms_index_list = [] #is it necessary?
        self.total_algs = len(self.algorithms)

        #build the algorithms_index_list and shuffle it
        self.numrounds = 10 
        for i in range(self.numrounds):
          self.algorithms_index_list.extend(range(self.total_algs)) 
        random.shuffle(self.algorithms_index_list)
#self.algorithms_index_list = [ 19, 4, 11, 12, 9, 3, 19, 16, 5, 22, 21, 6, 14, 21, 13, 4, 13, 0, 17, 21, 7, 9, 4, 10, 16, 8, 21, 7, 17, 11, 23, 23, 10, 11, 19, 17, 23, 6, 12, 0, 11, 15, 7, 16, 20, 3, 12, 16, 5, 18, 6, 17, 17, 23, 21, 8, 18, 18, 3, 18, 0, 6, 16, 15, 13, 16, 19, 19, 14, 2, 17, 20, 9, 11, 21, 23, 5, 13, 20, 7, 20, 10, 14, 6, 8, 9, 18, 6, 0, 13, 23, 3, 13, 7, 2, 14, 22, 21, 10, 20, 17, 8, 1, 0, 7, 4, 19, 9, 9, 20, 19, 1, 4, 5, 4, 19, 6, 1, 11, 15, 21, 18, 1, 15, 1, 3] 

        print(self.algorithms_index_list)
        self.numexperiments = len(self.algorithms_index_list) 
        ##########################################################

class InputResult(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes=['success'])

    def execute(self, userdata):
        keyserver = GetKeyStroke()
        c = -1
        print "Press [0-5] for clothes grabbed, 9 for error:"
        while (c < 0):
            c = keyserver()
            try:
                c = int(c)
            except ValueError:
                print "Error. Press [0-5] for clothes grabbed:"
                continue

        print "Your value was: " + str(c)
        pub = rospy.Publisher('/log/number_of_clothes', Int8, None, False, True)
        pub.publish(c)

        return 'success'

class PlaceHandler():
    def __init__(self):
        self.next_left_zone  = 1
        self.next_right_zone = 1

    def calculate_place_pose(self, pomdp_config):
        # If place_point is None it implies special basket
        if (pomdp_config.place_point == None):
            return self.get_basket_pose_st()

        #leave at the same side hard trick
        if (pomdp_config.grabbing_zone == pomdp_config.LEFT_ZONE):
            return self.get_left_place_pose()
        elif (pomdp_config.grabbing_zone == pomdp_config.RIGHT_ZONE):
            return self.get_right_place_pose()
        else:
            rospy.logerr("Value for graspping zone is not LEFT or RIGHT")
            return build_home_pose()

    def get_basket_pose_st(self):
        return build_pose_stamped_msg('/wam_link0', -0.165, 0.514, -0.188, 1, 0, 0, 0)

    def get_right_place_pose(self):
        if (self.next_right_zone == 1):
            pose = build_pose_stamped_msg('/wam_link0', 0.422, -0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_right_zone == 2):
            pose = build_pose_stamped_msg('/wam_link0',0.584, -0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_right_zone == 3):
            pose = build_pose_stamped_msg('/wam_link0',0.633, -0.160, -0.234, 1, 0, 0, 0)
        elif (self.next_right_zone == 4):
            pose = build_pose_stamped_msg('/wam_link0',0.476, -0.160, -0.234, 1, 0, 0, 0)
        else:
            rospy.logerr("Righ pose next zone is out of range: " + str(self.next_right_zone))

        self.next_right_zone = self.next_right_zone + 1

        # reset to 1 again
        if (self.next_right_zone > 4):
            self.next_right_zone = 1

        return pose

    def get_left_place_pose(self):
        if (self.next_left_zone == 1):
            pose = build_pose_stamped_msg('/wam_link0',0.422, 0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_left_zone == 2):
            pose = build_pose_stamped_msg('/wam_link0',0.584, 0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_left_zone == 3):
            pose = build_pose_stamped_msg('/wam_link0',0.633, 0.160, -0.234, 1, 0, 0, 0)
        elif (self.next_left_zone == 4):
            pose = build_pose_stamped_msg('/wam_link0',0.476, 0.160, -0.234, 1, 0, 0, 0)
        else:
            rospy.logerr("Left pose next zone is out of range: " + str(self.next_left_zone))

        self.next_left_zone = self.next_left_zone + 1

        # reset to 1 again
        if (self.next_left_zone > 4):
            self.next_left_zone = 1

        return pose

class FirstAction(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes    = ['ok','fail'],
                                   output_keys = ['pomdp_action','place_handler'])

    def execute(self, userdata):
        # Config and launch the deformable analysis

        randaction = random.randint(0,userdata.total_algs)
        print("initial random action: " % randaction)
        userdata.pomdp_action = randaction 
        pub = rospy.Publisher('/log/pomdp/action', Int8, None, False, True)
        pub.publish(response.action)

        if (not response):
            return 'fail'

        # Create the place handler
        userdata.place_handler = PlaceHandler()

        return 'ok'

class CountObjects(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes    = ['fail','success'],
                                   input_keys  = ['sm_pcl_table_left','sm_pcl_table_right'],
                                   output_keys = ['left_objects','right_objects'])
        self.number = -1
        # Subscribe to count numbers
        sub = rospy.Subscriber("/estirabot/skills/iri_textile_count/num_objects_found",
                               Int32, self.recieve_number)

    def execute(self, userdata):
        # Left objects
        self.calculate_objects(userdata.sm_pcl_table_left)
        userdata.left_objects = self.number
        pprint ("Left objects: " + str(self.number))
        pub = rospy.Publisher('/debug/observation/left_objects', Int8, None, False, True)
        pub.publish(self.number)

        #TODO what's this doing here?
        rospy.sleep(2)

        # Right objects
        self.calculate_objects(userdata.sm_pcl_table_right)
        userdata.right_objects = self.number
        pprint ("Right objects: " + str(self.number))
        pub = rospy.Publisher('/debug/observation/right_objects', Int8, None, False, True)
        pub.publish(self.number)

        return 'success'

    def recieve_number(self, data):
        self.number = data.data

    def calculate_objects(self, pcl):
        self.number = -1
        # Publish the pcl in /estirabot/skills/table/pcl
        pub = rospy.Publisher('/estirabot/skills/table/pcl', PointCloud2, None, False, True)
        pub.publish(pcl)
        # Wait to recieve the callback
        while (self.number == -1 and (not rospy.is_shutdown())):
            pprint("waiting")
            time.sleep(.1)

class PrePlaceCalculation(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes     = ['finish'],
                                   input_keys   = ['sm_pomdp_config'],
                                   output_keys  = ['pre_pose_st'])

    def execute(self, userdata):
        tf = Transform()
        tf.translation.z = -0.4

        pre_pose_st = PoseStamped()
        pre_pose_st.header.frame_id = '/wam_link0'
        pre_pose_st.header.stamp    = rospy.Time.now()
        pre_pose_st.pose = homogeneous_product_pose_transform(
                                        userdata.sm_pomdp_config.place_point.pose, tf)

        userdata.pre_pose_st = pre_pose_st
        return 'finish'

class GetNextAction(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes     = ['fail','finish','continue'],
                                   input_keys   = ['experiment_data'],
                                   output_keys  = ['pomdp_action', 'experiment_data'])


    def execute(self, userdata):

        ####################        
        pprint("iteration number = %d" % userdata.experiment_data.iteration_number)
        userdata.experiment_data.iteration_number += 1
        if (userdata.experiment_data.iteration_number == userdata.experiment_data.numexperiments):
            return 'finish'
        action = userdata.experiment_data.algorithms_index_list[userdata.experiment_data.iteration_number]
        ####################

        pub = rospy.Publisher('/log/pomdp/action', Int8, None, False, True)
        pub.publish(action) 
        userdata.pomdp_action = action
        return 'continue'

def get_sm_grasp_goal_callback(userdata, default_goal):
    return userdata.sm_grasp_goal

def get_sm_place_goal_callback(userdata, default_goal):
    goal = PlaceGoal()
    goal.place_locations.append(userdata.sm_pomdp_config.place_point)

    return goal

def construct_main_sm():
    sm = StateMachine(outcomes = ['succeeded','aborted'])

    factory_home = SM_ESTIRABOT_GoHome()
    sm_home = factory_home.build_sm()

    factory_analysis = SM_ESTIRABOT_DeformableAnalysis()
    sm_analysis = factory_analysis.build_sm()

    factory_grasp = SM_ESTIRABOT_GenerateGraspMsg()
    sm_grasp_generation = factory_grasp.build_sm()

    factory_move = SM_WAM_MoveFromPose('/estirabot/wam_wrapper/joints_move',
                                       '/estirabot/iri_wam_tcp_ik/get_wam_ik')
    sm_move = factory_move.build_sm()

    # Create the experiment_data
    sm.userdata.experiment_data = ExperimentData()
    # Create the place handler
    sm.userdata.place_handler = PlaceHandler()

    with sm:
        smach.StateMachine.add('START_ROUND', sm_home,
                          transitions={'success' : 'LOG_PCL',
                                       'aborted' : 'aborted'})
        smach.StateMachine.add('LOG_PCL', 
                                GetPCL("/estirabot/ceiling_kinect/camera/rgb/points"),
                transitions = {'success' : 'LOG_IMAGE',
                               'fail'    : 'aborted'},
                remapping   = {'pcl_RGB'  : 'sm_pcl_to_log'})
        smach.StateMachine.add('LOG_IMAGE',
                                TopicPublisher('/log/pcl_table', PointCloud2),
                          transitions={'finish': 'GET_NEXT_ACTION'},
                          remapping = {'msg'   : 'sm_pcl_to_log'})

        # 1 GET NEXT ACTION
        smach.StateMachine.add('GET_NEXT_ACTION', GetNextAction(),
                        transitions={'continue' : 'TRANSLATE_POMDP_ACTION',
                                     'finish'   : 'succeeded',
                                     'fail'     : 'aborted'})

        # 2. CALCULATE GRASP POINT
        smach.StateMachine.add('TRANSLATE_POMDP_ACTION', TranslatePomdp(),
            transitions = { 'image_from_left'  : 'GET_KINECT_IMAGE_FROM_LEFT',
                            'image_from_right' : 'GET_KINECT_IMAGE_FROM_RIGHT',
                            'fail':'aborted'},
            remapping   = { 'sm_pomdp_config' : 'sm_pomdp_config'})


        # 2.1.a GET KINECT IMAGE FROM LEFT
        smach.StateMachine.add('GET_KINECT_IMAGE_FROM_LEFT', 
                                GetPCL("/estirabot/skills/table/left/pcl"),
                transitions = {'success' : 'PERFORM_ANALISYS',
                               'fail'    : 'aborted'},
                remapping   = {'pcl_RGB'  : 'sm_pcl_output'})

        # 2.1.b GET KINECT IMAGE FROM RIGHT
        smach.StateMachine.add('GET_KINECT_IMAGE_FROM_RIGHT', 
                                GetPCL("/estirabot/skills/table/right/pcl"),
                transitions = {'success' : 'PERFORM_ANALISYS',
                               'fail'    : 'aborted'},
                remapping   = {'pcl_RGB'  : 'sm_pcl_output'})

        # 2.2 ANALYSIS
        smach.StateMachine.add('PERFORM_ANALISYS', sm_analysis,
                transitions = {'success'        : 'CALCULATE_GRASP',
                               'timeout'        : 'PERFORM_ANALISYS', # hack relaunch since descriptors died
                               'fail'           : 'aborted'},
                remapping   = {'deformable_config' : 'sm_pomdp_config', # compatibles
                               'pcl_RGB'           : 'sm_pcl_output',
                               'best_grasp_pose'   : 'sm_best_grasp_pose'})
        # 2.3 CALCULATE GRASP
        smach.StateMachine.add('CALCULATE_GRASP', sm_grasp_generation,
                transitions = {'success'        : 'DO_GRASP',
                               'no_ik_solution' : 'CALCULATE_GRASP',
                               'fail'           : 'aborted',
                               'all_ik_fail'    : 'GET_TABLE_IMAGE_FROM_LEFT'},
                remapping   = {'pomdp_grasp_config'  : 'sm_pomdp_config', # compatibles
                               'grasp_pose'          : 'sm_best_grasp_pose',
                               'grasp_goal'          : 'sm_grasp_goal'})
        # 3.1 GRASP
        smach.StateMachine.add('DO_GRASP',
                              smach_ros.SimpleActionState('/estirabot/skills/grasp/pickup',
                                                          PickupAction,
                                                          goal_cb = get_sm_grasp_goal_callback,
                                                          input_keys=['sm_grasp_goal']),
                               {'succeeded': 'DO_LIFT',
                                'preempted': 'aborted',
                                'aborted'  : 'aborted'})
        # 3.2 LIFT
        smach.StateMachine.add('DO_LIFT', sm_home,
                        transitions={'success' : 'CHECK_GRASP',
                                     'aborted' : 'aborted'})
        # 3.2b CHECK
        smach.StateMachine.add('CHECK_GRASP', InputResult(),
                        transitions={'success' : 'CALCULATE_PRE_PLACE'})

        # 3.3.1 DO PRE PLACE
        smach.StateMachine.add('CALCULATE_PRE_PLACE', PrePlaceCalculation(),
                        transitions={'finish' : 'MOVE_TO_PRE_PLACE'},
                          remapping={'pre_pose_st':'sm_pre_pose_st'})

        smach.StateMachine.add('MOVE_TO_PRE_PLACE', sm_move,
                        transitions={'success' : 'DO_PLACE',
                                     'aborted' : 'aborted',
                                     'no_kinematic_solution'  : 'DO_PLACE'}, #TODO recompute pre place point 
                          remapping={'pose_st':'sm_pre_pose_st'})

        # 3.3.2 DO PLACE
        smach.StateMachine.add('DO_PLACE',
                       smach_ros.SimpleActionState('/estirabot/skills/iri_place/place',
                                                  PlaceAction,
                                                  goal_cb = get_sm_place_goal_callback,
                                                  input_keys=['sm_pomdp_config']),
                       {'succeeded': 'BACK_TO_HOME',
                        'preempted': 'aborted',
                        'aborted'  : 'aborted'})

        smach.StateMachine.add('BACK_TO_HOME', sm_home,
                          transitions={'success' : 'GET_TABLE_IMAGE_FROM_LEFT',
                                       'aborted' : 'aborted'})

        # 4. OBSERVATION
        # 4.1 GET PCL FROM LEFT
        smach.StateMachine.add('GET_TABLE_IMAGE_FROM_LEFT', 
                      GetPCL("/estirabot/skills/table/left/pcl"),
                transitions = {'success' : 'GET_TABLE_IMAGE_FROM_RIGHT',
                               'fail'    : 'aborted'},
                remapping   = {'pcl_RGB'  : 'sm_pcl_table_left'})
        # 4.2 GET PCL FROM RIGHT
        smach.StateMachine.add('GET_TABLE_IMAGE_FROM_RIGHT', 
                      GetPCL("/estirabot/skills/table/right/pcl"),
                transitions = {'success' : 'COUNT_OBJECTS',
                               'fail'    : 'aborted'},
                remapping   = {'pcl_RGB'  : 'sm_pcl_table_right'})
        # 4.3 GET NUMBER OF OBJECTS
        smach.StateMachine.add('COUNT_OBJECTS', CountObjects(),
                        transitions={'success' : 'START_ROUND',
                                     'fail'    : 'aborted'})

    return sm

def main():
    rospy.init_node("pomdp_sm")
    sm_main = construct_main_sm()

    # Run state machine introspection server for smach viewer
    sis = smach_ros.IntrospectionServer('estirabot_smach_replicable', sm_main,
                                        '/estirabot_smach_replicable')

    sis.start()

    outcome = sm_main.execute()

    rospy.spin()
    sis.stop()

if __name__ == "__main__":
    main()