pomdp_execution.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 the parent directory for common
import os, sys
_root_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.insert(0, _root_dir)

from estirabot_common import *
from pomdp_create_model import *

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 estirabot_msgs.srv import obs2action, obs2actionRequest
from geometry_msgs.msg import Transform, PoseStamped, Pose
from pprint import pprint
from std_msgs.msg import Int8, Int32

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()

        if (pomdp_config.grabbing_zone == pomdp_config.LEFT_ZONE):
            return self.get_right_place_pose()
        elif (pomdp_config.grabbing_zone == pomdp_config.RIGHT_ZONE):
            return self.get_left_place_pose()
        else:
            rospy.logerr("Value for graspping zone is not LEFT or RIGHT")
            return build_home_pose()

    def get_basket_pose(self):
        return build_pose(-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(0.422, -0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_right_zone == 2):
            pose = build_pose(0.584, -0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_right_zone == 3):
            pose = build_pose(0.633, -0.160, -0.234, 1, 0, 0, 0)
        elif (self.next_right_zone == 4):
            pose = build_pose(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(0.422, 0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_left_zone == 2):
            pose = build_pose(0.584, 0.390, -0.234, 1, 0, 0, 0)
        elif (self.next_left_zone == 3):
            pose = build_pose(0.633, 0.160, -0.234, 1, 0, 0, 0)
        elif (self.next_left_zone == 4):
            pose = build_pose(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

def get_config_from_pomdp_action(pomdp_action):

    if (pomdp_action > 21):
        rospy.logfatal("Pomdp action id out of range: " + str(pomdp_action))

    server = PomdpConfigFactory()

    if  (pomdp_action == 20):
        magic_number = 114
    elif (pomdp_action == 21):
        magic_number = 114
    elif (pomdp_action % 10 == 0):
        magic_number = 1
    elif (pomdp_action % 10 == 1):
        magic_number = 3
    elif (pomdp_action % 10 == 2):
        magic_number = 5
    elif (pomdp_action % 10 == 3):
        magic_number = 11
    elif (pomdp_action % 10 == 4):
        magic_number = 14
    elif (pomdp_action % 10 == 5):
        magic_number = 101
    elif (pomdp_action % 10 == 6):
        magic_number = 103
    elif (pomdp_action % 10 == 7):
        magic_number = 105
    elif (pomdp_action % 10 == 8):
        magic_number = 111
    elif (pomdp_action % 10 == 9):
        magic_number = 114

    # From magic number get configuration
    config = server.get_instance(magic_number)

    # Also set the grabbing zone
    if (pomdp_action == 20):
        config.grabbing_zone = config.RIGHT_ZONE
    elif (pomdp_action == 21):
        config.grabbing_zone = config.LEFT_ZONE
    elif (pomdp_action > 9):
        config.grabbing_zone = config.LEFT_ZONE
    else:
        config.grabbing_zone = config.RIGHT_ZONE

    # 20 and 21 should get an special value for place.
    # Hack for this is set place_point to None and check for it
    # in PlaceHandler. Dirty as hell.
    if (pomdp_action == 20) or (pomdp_action == 21):
        config.place_point = None

    return config

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
        request = obs2actionRequest()
        request.first       = True
        handler = rospy.ServiceProxy('/estirabot/skills/planner/obs2action', obs2action)
        response = handler(request)

        pprint(response)
        userdata.pomdp_action = response.action
        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 LogPCL(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes=['success','fail'],
                                   input_keys=['sm_pcl_to_log'],
                                   output_keys=[''])

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


class TranslatePomdp(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes    = ['fail','image_from_left','image_from_right'],
                                   input_keys  = ['pomdp_action','place_handler'],
                                   output_keys = ['sm_pomdp_config','place_handler'])

    def execute(self, userdata):
        # Stupid trick to not declare sm_pomdp_config as input (SMACH goodies)
        config                    = get_config_from_pomdp_action(userdata.pomdp_action)
        config.place_point        = userdata.place_handler.calculate_place_pose(config)

        userdata.sm_pomdp_config = config

        pub_config = rospy.Publisher('/debug/pomdp_config', PomdpGraspConfig, None, False, True)
        pub_config.publish(config)
        pprint(config)

        if (config.grabbing_zone == config.LEFT_ZONE):
            return 'image_from_left'
        else:
            return 'image_from_right'

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)

        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 PrePlace(smach.State):
    def __init__(self):
        smach.State.__init__(self, outcomes     = ['fail','success'],
                                   input_keys   = ['sm_pomdp_config'])

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

        return wam_cartesian_move(homogeneous_product_pose_transform(userdata.sm_pomdp_config.place_point, 
                                                                     tf))


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

    def execute(self, userdata):
        request = obs2actionRequest()
        # Build observation
        right_objs = userdata.right_objects
        left_objs  = userdata.left_objects

        if (right_objs > 3):
            right_objs = 4

        if (left_objs > 4):
            left_objs = 4
        
        request.first       = False
        request.observation = right_objs * 5 + left_objs
        pub = rospy.Publisher('/log/pomdp/observation', Int8, None, False, True)
        pub.publish(request.observation)

        pprint(request)
        handler = rospy.ServiceProxy('/estirabot/skills/planner/obs2action', obs2action)
        response = handler(request)

        if (not response):
            return 'fail'

        pprint(response)
        if (response.goal_reached == True):
            return 'finish'

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

        return 'continue'


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

    with sm:
        smach.StateMachine.add('GIVE_FIRST_ACTION', FirstAction(),
                transitions = { 'ok':'START_ROUND',
                                'fail':'aborted'},
                remapping   = { 'pomdp_action' : 'pomdp_action'})

        smach.StateMachine.add('START_ROUND', EstirabotGoHome(),
                          transitions={'success': 'LOG_PCL',
                                       'fail'   : 'aborted'})
        smach.StateMachine.add('LOG_PCL', 
                                EstirabotGetPCL("/estirabot/ceiling_kinect/camera/rgb/points"),
                transitions = {'success' : 'LOG_IMAGE',
                               'fail'    : 'aborted'},
                remapping   = {'pcl_RGB'  : 'sm_pcl_to_log'})
        smach.StateMachine.add('LOG_IMAGE', LogPCL(),
                          transitions={'success': 'TRANSLATE_POMDP_ACTION',
                                       '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', 
                                EstirabotGetPCL("/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', 
                              EstirabotGetPCL("/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', PerformAnalysis(),
                transitions = {'success'        : 'CALCULATE_GRASP',
                               'timeout'        : 'PERFORM_ANALISYS', # hack relaunch since descriptors died
                               'fail'           : 'aborted'},
                remapping   = {'pcl_RGB'         : 'sm_pcl_output',
                               'best_grasp_pose' : 'sm_best_grasp_pose'})
        # 2.3 CALCULATE GRASP
        smach.StateMachine.add('CALCULATE_GRASP', CalculateGrasp(),
                transitions = {'success'        : 'DO_GRASP',
                               'no_ik_solution' : 'CALCULATE_GRASP',
                               'fail'           : 'aborted',
                               'all_ik_fail'    : 'GET_TABLE_IMAGE_FROM_LEFT'},
                remapping   = {'pcl_RGB'         : 'sm_pcl_output',
                               'best_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', EstirabotGoHome(),
                        transitions={'success' : 'DO_PRE_PLACE',
                                     'fail'    : 'aborted'})
        # 3.3.1 DO PRE PLACE
        smach.StateMachine.add('DO_PRE_PLACE', PrePlace(),
                        transitions={'success' : 'DO_PLACE',
                                     'fail'    : 'aborted'})
        #
        # 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', EstirabotGoHome(),
                          transitions={'success': 'GET_TABLE_IMAGE_FROM_LEFT',
                                       'fail'   : 'aborted'})


        # 4. OBSERVATION
        # 4.1 GET PCL FROM LEFT
        smach.StateMachine.add('GET_TABLE_IMAGE_FROM_LEFT', 
                      EstirabotGetPCL("/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', 
                      EstirabotGetPCL("/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' : 'UPDATE_BELIEVE_STATE',
                                     'fail'    : 'aborted'})
        # 5 UPDATE BELIEVE STATE
        smach.StateMachine.add('UPDATE_BELIEVE_STATE', UpdateBelieve(),
                        transitions={'continue' : 'START_ROUND',
                                     'finish'   : 'succeeded',
                                     '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()