pomdp_create_model.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

from smach import StateMachine, CBState, Iterator
from smach_ros import ServiceState
from actionlib import *
from actionlib.msg import *
from std_msgs.msg import String
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 geometry_msgs.msg import Transform, PoseStamped, Pose
from pprint import pprint
from std_msgs.msg import Int8

from iri_common_smach.st_get_pcl          import GetPCL
from iri_common_smach.utils_msg           import build_pose_stamped_msg, build_transform_msg
from iri_common_smach.get_keystroke       import GetKeyStroke

from estirabot_apps_base.utils_pomdp            import PomdpConfigFactory
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

class PublishGraspID(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes=['success','fail'],
                             input_keys=['sm_pomdp_config'])

    def execute(self, userdata):
        print "Publish grasp id"
        #  pprint(userdata.sm_pomdp_config)
#raw_input()
        return 'success'

class HandlePomdpConfig(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes=['ok'], input_keys=['algorithms','algorithms_index'],
                                                    output_keys=['sm_pomdp_config','algorithms_index'])

    def execute(self, userdata):
        magic_number = userdata.algorithms[userdata.algorithms_index]

        factory = PomdpConfigFactory()
        userdata.sm_pomdp_config = factory.get_instance(magic_number)
        # Publish the magic number
        pub = rospy.Publisher('/log/config_magic_number', Int8, None, False, True)
        pub.publish(magic_number)
        pprint("MAGIC NUMBER: " + str(magic_number))

        userdata.algorithms_index = userdata.algorithms_index + 1
        pprint("INDEX: " + str(userdata.algorithms_index))

        return 'ok'

#TODO Ducplicated class @estirabot_apps_base
class PerformAnalysis(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes=['success','fail','timeout'],
                                   input_keys=['pcl_RGB','sm_pomdp_config'],
                                   output_keys=['best_grasp_pose'])

    def execute(self, userdata):
        # Log the PCL
        pub = rospy.Publisher('/log/pcl_table_scene', PointCloud2,  None, False, True)
        #pub.publish(userdata.pcl_RGB)

        # Config and launch the deformable analysis
        handler = rospy.ServiceProxy('/estirabot/skills/deformable_analysis/do_deformable_analysis',
                                     DoDeformableAnalysis)

        analysis_request = DoDeformableAnalysisRequest()
        analysis_request.pcl_to_analyze = userdata.pcl_RGB
        analysis_request.fusion_criteria_id = userdata.sm_pomdp_config.best_pose_algorithm_id

        try:
            response = handler(analysis_request)

            if (not response):
                return 'fail'

        except rospy.ServiceException, e:
            return 'timeout'

        pose_st = response.deformable_analysis.graspability.best_grasp_pose
        userdata.best_grasp_pose = pose_st
        print("end analysis")

        return 'success'

class CalculateGrasp(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes=['success','fail','no_ik_solution','all_ik_fail'],
                                   input_keys=['best_grasp_pose','sm_pomdp_config'],
                                   output_keys=['sm_pomdp_config','grasp_goal'])
        self.tf_manager = TransformManager()

    def execute(self, userdata):
        # shorcuts
        pomdp_config  = userdata.sm_pomdp_config
        grasp_pose_st = userdata.best_grasp_pose

        rospy.set_param('/estirabot/skills/grasp/fingers_configuration', pomdp_config.fingers_grasp_configs)

        # grasp_pose_st comes referecence to camera frame_id
        # we need to transform the message to wam base frame wam_link0 
        grasp_pose_st_wam = self.tf_manager.transform_pose('wam_link0', grasp_pose_st)

        # Orientation is fixed so ignoring previous
        grasp_pose_st_wam.pose.orientation.x = 0
        grasp_pose_st_wam.pose.orientation.y = 0
        grasp_pose_st_wam.pose.orientation.z = 0
        grasp_pose_st_wam.pose.orientation.w = 1.0

        pprint("Grasp original")
        pub = rospy.Publisher('/debug/grasp/original', PoseStamped, None, False, True)
        pub.publish(grasp_pose_st_wam)

        # real grasp is build applying grasp modifier to grasp_pose
        current_modifier     = pomdp_config.approach_config.grasp_modifier_used
        pprint("Using grasp modifier number: " + str(current_modifier))
        real_grasp_st        = PoseStamped()
        real_grasp_st.header = grasp_pose_st_wam.header
        real_grasp_st.pose   = homogeneous_product_pose_transform(grasp_pose_st_wam.pose,
                                                          pomdp_config.approach_config.grasp_modifiers[current_modifier])
        if (real_grasp_st.pose.position.z < -0.282):
            real_grasp_st.pose.position.z = -0.282

        pprint("Grasp modifier applied")
        pub = rospy.Publisher('/debug/grasp/real', PoseStamped, None, False, True)
        pub.publish(real_grasp_st)

        # pre-grasp posture is relative to real_grasp pose
        pre_grasp_pose_st        = PoseStamped()
        pre_grasp_pose_st.header = grasp_pose_st_wam.header

        pre_grasp_pose_st.pose = homogeneous_product_pose_transform(real_grasp_st.pose,
                                                          pomdp_config.approach_config.pre_grasp_modifier)
        pprint("PRE-Grasp modifier applied")
        pub = rospy.Publisher('/debug/grasp/pre_grasp', PoseStamped, None, False, True)
        pub.publish(pre_grasp_pose_st)

        # GRASP is Kinematic compliant so save and log the final configuration
        pub_config = rospy.Publisher('/log/pomdp_config', PomdpGraspConfig, None, False, True)
        pub_config.publish(pomdp_config)

        # For both grasp and pre-grasp we need the inverse kinematics and get joinst positions
        try:
            joints_grasp     = get_inverse_kinematics_from_pose(real_grasp_st.pose)
            joints_pre_grasp = get_inverse_kinematics_from_pose(pre_grasp_pose_st.pose)

        #except NoKinematicSolutionException, e:
        except rospy.ServiceException, e:
            # Let's go to next grasp_modifier (if any)
            if (current_modifier < len(pomdp_config.approach_config.grasp_modifiers)-1):
                # Mark next configuration to use 
                pomdp_config.approach_config.grasp_modifier_used = pomdp_config.approach_config.grasp_modifier_used + 1
                pomdp_config.approach_config.pre_grasp_modifier
                return 'no_ik_solution'
            else:
                # all modifiers were tested and are not valid. Return error
                pub = rospy.Publisher('/log/number_of_clothes', Int8, None, False, True)
                pub.publish(0)
                rospy.logerr("Inverse Kinematics has no solution")
                return 'all_ik_fail'

        #     rospy.logerr("WamIK service failed")
        #           return 'fail'

        userdata.grasp_goal =  build_simple_pickup_goal(joints_grasp,
                                                        joints_pre_grasp,
                                                        pomdp_config.fingers_grasp_configs)
        return 'success'

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

def get_sm_place_goal_callback(userdata, default_goal):
    pose_stamped                 = PoseStamped()
    pose_stamped.header.frame_id = 'wam_link0'
    pose_stamped.header.stamp    = rospy.Time.now()
    pose_stamped.pose            = userdata.sm_pomdp_config.place_point

    goal = PlaceGoal()
    goal.place_locations.append(userdata.sm_pomdp_config.place_point)

    return goal

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

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

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

        return 'ok'

# Iterator over rounds given a fixed config
def construct_round_iterator_sm():

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

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

    factor_grasp = SM_ESTIRABOT_GenerateGraspMsg()
    sm_grasp_generation = factor_grasp.build_sm()

    # 1. START_ROUND state machine
    sm_start_round = smach.StateMachine(outcomes=['success','aborted'],
                                        input_keys=['sm_pomdp_config'])

    with sm_start_round:
        smach.StateMachine.add('PUBLISH_GRASP_ID', PublishGraspID(),
                          transitions={'success'   : 'MOVE_TO_HOME_POSITION',
                                       'fail'      : 'aborted'})


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

    # 2. CALCULATE_GRASP_POINT state machine
    sm_calculate_grasp_point = smach.StateMachine(outcomes=['success','aborted','all_ik_fail'],
                                                  input_keys=['sm_pomdp_config'],
                                                  output_keys=['sm_grasp_goal'])

    with sm_calculate_grasp_point:
        # Get kinect image and process and make analisys
        smach.StateMachine.add('GET_KINECT_IMAGE', 
                               GetPCL("/estirabot/skills/pcl_distance_filtered"),
                transitions = {'success' : 'PERFORM_ANALISYS',
                               'fail'    : 'aborted'},
                remapping   = {'pcl_RGB'         : 'sm_pcl_output'})
        smach.StateMachine.add('PERFORM_ANALISYS', sm_analysis,
                transitions = {'success'        : 'CALCULATE_GRASP',
                               'fail'           : 'aborted',
                               'timeout'        : 'aborted'},
                remapping   = {'deformable_config' : 'sm_pomdp_config', # compatibles
                               'pcl_RGB'           : 'sm_pcl_output',
                               'best_grasp_pose'   : 'sm_best_grasp_pose'})
        smach.StateMachine.add('CALCULATE_GRASP', sm_grasp_generation,
                transitions = {'success'        : 'success',
                               'fail'           : 'aborted',
                               'no_ik_solution' : 'CALCULATE_GRASP',
                               'all_ik_fail'    : 'all_ik_fail'},
                remapping   = {'grasp_pose'         : 'sm_best_grasp_pose',
                               'pomdp_grasp_config' : 'sm_pomdp_config',
                               'grasp_goal'         : 'sm_grasp_goal'})

    # 3. PICK & LIFT
    sm_pick_and_lift = smach.StateMachine(outcomes=['success','aborted'],
                                          input_keys=['sm_grasp_goal'])

    with sm_pick_and_lift:
        # TODO: endless loop if fail
        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'})

        smach.StateMachine.add('DO_LIFT', sm_home,
                        transitions={'success' : 'success',
                                     'aborted' : 'aborted'})
    # 4. CHECK
    sm_check_result = smach.StateMachine(outcomes=['success'])

    with sm_check_result:
        smach.StateMachine.add('INPUT_RESULT', InputResult(),
                transitions = { 'ok' : 'success'})

    # 5. DROP
    sm_drop = smach.StateMachine(outcomes=['success','aborted'],
                                   input_keys=['sm_pomdp_config'])

    with sm_drop:
        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': 'success',
                        'preempted': 'aborted',
                        'aborted'  : 'aborted'})

    # MAIN STATE MACHINE
    sm = StateMachine(outcomes = ['continue','aborted'],
                      input_keys = ['algorithms','algorithms_index'],
                      output_keys = ['algorithms','algorithms_index'])
    sm.userdata.rounds = 10

    with sm:
        smach.StateMachine.add('CHOOSE_POMDP_CONFIG', HandlePomdpConfig(),
            transitions = {'ok':'ROUND_ITERATOR'})

        round_iterator = Iterator(outcomes = ['succeeded','aborted'],
                               input_keys = ['sm_pomdp_config'],
                               it = range(0, sm.userdata.rounds),
                               output_keys = ['sm_pomdp_config'],
                               it_label = 'index',
                               exhausted_outcome = 'succeeded')
        with round_iterator:
            container_sm = StateMachine(outcomes = ['succeeded','aborted','continue'],
                                        input_keys = ['sm_pomdp_config'],
                                        output_keys = ['sm_pomdp_config'])
            with container_sm:
                # 1 START ROUND
                smach.StateMachine.add('START_ROUND', sm_start_round,
                        transitions={'success': 'CALCULATE_GRASP_POINT',
                                     'aborted': 'aborted'})
                # 2 CALCULATE GRASP POINT
                smach.StateMachine.add('CALCULATE_GRASP_POINT', sm_calculate_grasp_point,
                        transitions={'success'     : 'PICK_AND_LIFT',
                                     'aborted'     : 'aborted',
                                     'all_ik_fail' : 'continue'})

                # 3 PICK AND LIFT
                smach.StateMachine.add('PICK_AND_LIFT', sm_pick_and_lift,
                        transitions={'success': 'CHECK_RESULT',
                                     'aborted': 'aborted' })
                # 4. CHECK
                smach.StateMachine.add('CHECK_RESULT', sm_check_result,
                        transitions = {'success': 'DROP'})

                # 5. DROP
                smach.StateMachine.add('DROP', sm_drop,
                        transitions = {'success': 'continue',
                                       'aborted': 'aborted'})

            #close container_sm
            Iterator.set_contained_state('CONTAINER_STATE', 
                                         container_sm, 
                                         loop_outcomes=['continue'])
        #close the round_iterator
        StateMachine.add('ROUND_ITERATOR',round_iterator,
                     {'succeeded':'continue',
                      'aborted':'aborted'})
    return sm

# Main iterator to choose algorithms and run each of them
# several times
def construct_alg_iterator_sm():
    sm = StateMachine(outcomes = ['succeeded','aborted'], output_keys = ['algorithms_index','algorithms'])
    sm.userdata.algorithms_index = 0
    sm.userdata.algorithms = [ 4 ]

    sm.userdata.total_algs = len(sm.userdata.algorithms)

    sm_container = construct_round_iterator_sm()

    with sm:
        alg_iterator = Iterator(outcomes = ['succeeded','aborted'],
                               input_keys = ['algorithms','algorithms_index'],
                               it = range(0, sm.userdata.total_algs),
                               output_keys = ['algorithms','algorithms_index'],
                               it_label = 'index',
                               exhausted_outcome = 'succeeded')

        with alg_iterator:
            Iterator.set_contained_state('ALG_CONTAINER_STATE',
                                    sm_container,
                                    loop_outcomes=['continue'])

        smach.StateMachine.add('ALGORITHM_ITERATOR', alg_iterator,
                 {'succeeded':'succeeded',
                  'aborted':'aborted'})

    return sm

def main():
    rospy.init_node("pomdp_sm")
    sm_iterator = construct_alg_iterator_sm()

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

    sis.start()


    outcome = sm_iterator.execute()

    rospy.spin()
    sis.stop()

if __name__ == "__main__":
    main()