sm_cleaning_learning.py

Go to the documentation of this file.

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import roslib; roslib.load_manifest('iri_clean_board')
import rospy
import smach
import smach_ros

from smach_ros import SimpleActionState
from smach_ros import ServiceState

from actionlib import *
from actionlib.msg import *

from iri_common_smach import st_sleep

from iri_clean_board.srv import *
from estirabot_msgs.msg import *
from estirabot_msgs.srv import *
from iri_planning_msgs.srv import *

class NeedInfo(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['next','failed'])
                self.counter = 0

        def execute(self, userdata):
                rospy.logdebug('Executing state NeedInfo')
                return 'next'


class GetInfoCamera(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['succeeded','preempted','aborted'],
                                     input_keys=[],
                                     output_keys=['segmented_image', 'pointcloud', 'plane_coefficients'])
                self.service_topic = '/iri_clean_perception/get_board_information/get_board_info_srv'
        
        def execute(self, userdata):
                rospy.logdebug('Executing state GetInfoCamera')
                start_timer = time.time()
                
                rospy.wait_for_service(self.service_topic)
                try:
                        get_service = rospy.ServiceProxy(self.service_topic, BoardInfo)
                        resp = get_service()
                        
                        userdata.segmented_image = resp.segmented_image
                        userdata.pointcloud = resp.pointcloud
                        userdata.plane_coefficients = resp.plane_coefficients
                        
                        stop_timer = time.time()
                        global time_perception
                        time_perception = time_perception + (stop_timer - start_timer)
                        return 'succeeded'
                
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'aborted'


class PradaPlanning(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['succeeded','preempted','aborted','no_plan'],
                                     input_keys=['rules_updated_sm'],
                                     output_keys=['target_ellipses','actions'])
                self.service_topic = '/iri_clean_plan/prada_planner/getBestAction'
        
        def execute(self, userdata):
                rospy.logdebug('Executing state PradaPlanning')
                
                start_timer = time.time()
                rospy.wait_for_service(self.service_topic)
                try:
                        get_cleaning_plan = rospy.ServiceProxy(self.service_topic, PradaPlanSrv)
                        resp = get_cleaning_plan(userdata.rules_updated_sm)
                        
                        stop_timer = time.time()
                        global time_planning
                        time_planning = time_planning + (stop_timer - start_timer)
                        
                        if resp.succeeded:
                                userdata.target_ellipses = resp.target_ellipses
                                userdata.actions = resp.actions
                                return 'succeeded'
                        else:
                                return 'no_plan'
                
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'aborted'


class CalculateTrajectories(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['succeeded','preempted','aborted','no_movements'],
                                     input_keys=['target_ellipses_sm','actions_sm','ellipses_sm','segmented_pointcloud_sm','plane_coefficients_sm'],
                                     output_keys=['poses','poses_2d_coordinates','secondary_arm'])
                self.service_topic = '/iri_clean_plan/cleaning_plan/get_trajectories_from_plan'
        
        def execute(self, userdata):
                rospy.logdebug('Executing state CalculateTrajectories')
                start_timer = time.time()
                rospy.wait_for_service(self.service_topic)
                try:
                        get_cleaning_plan = rospy.ServiceProxy(self.service_topic, PlanToTrajectories)
                        resp = get_cleaning_plan(userdata.segmented_pointcloud_sm, userdata.plane_coefficients_sm, userdata.target_ellipses_sm, userdata.actions_sm)
                        
                        stop_timer = time.time()
                        global time_state
                        time_state = time_state + (stop_timer - start_timer)
                        
                        if resp.movement:
                                userdata.poses = resp.poses
                                userdata.poses_2d_coordinates = resp.poses_2d_coordinates
                                userdata.secondary_arm = resp.secondary_arm
                                return 'succeeded'
                        else:
                                return 'no_movements'
                
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'aborted'


class GetStateRepresentation(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['succeeded','preempted','aborted','is_clean','get_full_info'],
                                     input_keys=['segmented_image_sm','segmented_pointcloud_sm','plane_coefficients_sm'],
                                     output_keys=['dirty_areas','distances','state_string','arm_position','segmented_pointcloud_sm','plane_coefficients_sm'])
                self.service_topic = '/iri_clean_plan/cleaning_plan/get_dirty_areas_state_representation'
                self.surface_clean_perceptions = 0
        
        def execute(self, userdata):
                rospy.logdebug('Executing state GetStateRepresentation')
                start_timer = time.time()
                rospy.wait_for_service(self.service_topic)
                try:
                        get_state = rospy.ServiceProxy(self.service_topic, StateRepresentation)
                        resp = get_state(userdata.segmented_image_sm)
                        
                        stop_timer = time.time()
                        global time_state
                        time_state = time_state + (stop_timer - start_timer)
                        
                        if resp.is_clean:
                                if self.surface_clean_perceptions < 2:
                                        self.surface_clean_perceptions += 1
                                        return 'get_full_info'
                                else:
                                        return 'is_clean'
                        else:
                                self.surface_clean_perceptions = 0
                                userdata.dirty_areas = resp.dirty_areas
                                userdata.distances = resp.distances
                                userdata.arm_position = resp.arm_position
                                userdata.state_string = resp.state_string
                                return 'succeeded'
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'aborted'


class LearnerAddAction(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['needs_learning','succeeded','preempted','aborted'],
                                     input_keys=['target_ellipses_sm','actions_sm','dirty_areas_sm','distances_sm','state_string_sm','movements_ok_sm'],
                                     output_keys=[])
                self.service_topic = '/iri_learning/rule_learner/add_action'
        
        def execute(self, userdata):
                rospy.logdebug('Executing state LearnerAddAction')
                start_timer = time.time()
                rospy.wait_for_service(self.service_topic)
                try:
                        get_cleaning_plan = rospy.ServiceProxy(self.service_topic, AddAction)
                        resp = get_cleaning_plan(userdata.actions_sm[0], userdata.dirty_areas_sm, userdata.target_ellipses_sm[0], userdata.distances_sm, userdata.state_string_sm, userdata.movements_ok_sm)
                        
                        stop_timer = time.time()
                        global time_learning
                        time_learning = time_learning + (stop_timer - start_timer)
                        
                        if resp.needs_learning:
                                return 'needs_learning'
                        else:
                                return 'succeeded'
                
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'aborted'


class LearnerAddState(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['succeeded','preempted','aborted'],
                                     input_keys=['segmented_image_sm'],
                                     output_keys=['rules_updated'])
                self.service_topic = '/iri_learning/rule_learner/add_state'
        
        def execute(self, userdata):
                rospy.logdebug('Executing state LearnerAddState')
                start_timer = time.time()
                rospy.wait_for_service(self.service_topic)
                try:
                        get_cleaning_plan = rospy.ServiceProxy(self.service_topic, AddState)
                        resp = get_cleaning_plan(userdata.segmented_image_sm)
                        
                        stop_timer = time.time()
                        global time_learning
                        time_learning = time_learning + (stop_timer - start_timer)
                        
                        userdata.rules_updated = resp.rules_updated
                        return 'succeeded'
                
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'aborted'


class MoveArm(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['succeeded','failed'],
                                        input_keys=['poses','plane_coefficients_sm','poses_2d_coordinates','secondary_arm'],
                                        output_keys=['movements_ok'])

        def execute(self, userdata):
                rospy.logdebug('Waiting for service')
                start_timer = time.time()
                rospy.wait_for_service('/iri_clean_plan/arm_movements_by_pose/clean_movement')
                try:
                        move_arm = rospy.ServiceProxy('/iri_clean_plan/arm_movements_by_pose/clean_movement', ArmMovementsPosesSrv)
                        resp = move_arm(userdata.poses, userdata.secondary_arm, userdata.plane_coefficients_sm, 1)
                        
                        stop_timer = time.time()
                        global time_moving
                        time_moving = time_moving + (stop_timer - start_timer)
                        
                        if resp.success:
                                print 'Executing service success'
                                userdata.movements_ok = True
                                return 'succeeded'
                        else:
                                # publish unreachable points
                                pub = rospy.Publisher('/iri_clean_perception/get_board_information/unreachable_position', ImagePoint)
                                for idx in resp.too_far_points_indexes:
                                        pub.publish(userdata.poses_2d_coordinates[idx])
                                print 'Executing service failed. Maybe IK.'
                                userdata.movements_ok = False
                                return 'succeeded'
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'failed'


class MoveArmBase(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['next','failed'])

        def execute(self, userdata):
                rospy.loginfo('Waiting for service')
                start_timer = time.time()
                rospy.wait_for_service('/iri_clean_plan/arm_movements_by_pose/clean_movement')
                try:
                        # make a empty service message
                        clean_mov = ArmMovementsPosesSrv()
                        clean_mov.movement_type = 0
                        clean_mov.poses = [];
                        move_arm = rospy.ServiceProxy('/iri_clean_plan/arm_movements_by_pose/clean_movement', ArmMovementsPosesSrv)
                        resp = move_arm(clean_mov.poses, [], [], 0)
                        
                        stop_timer = time.time()
                        global time_moving
                        time_moving = time_moving + (stop_timer - start_timer)
                        
                        if resp.success:
                                print 'Executing service success'
                                return 'next'
                        else:
                                print 'Executing service  failed'
                                return 'next'
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'failed'


class MoveArmBaseWait(smach.State):
        def __init__(self):
                smach.State.__init__(self, outcomes=['next','failed'])

        def execute(self, userdata):
                rospy.loginfo('Waiting for service')
                start_timer = time.time()
                rospy.wait_for_service('/iri_clean_plan/arm_movements_by_pose/clean_movement')
                try:
                        # make a empty service message
                        clean_mov = ArmMovementsPosesSrv()
                        clean_mov.movement_type = 0
                        clean_mov.poses = [];
                        move_arm = rospy.ServiceProxy('/iri_clean_plan/arm_movements_by_pose/clean_movement', ArmMovementsPosesSrv)
                        resp = move_arm(clean_mov.poses, [], [], 0)
                        
                        stop_timer = time.time()
                        global time_extra_learn
                        time_extra_learn = time_extra_learn + (stop_timer - start_timer) + 1.0
                        
                        if resp.success:
                                print 'Executing service success'
                                return 'next'
                        else:
                                print 'Executing service  failed'
                                return 'next'
                except rospy.ServiceException, e:
                        print "Service call failed: %s"%e
                        return 'failed'


def calculate_trajectories_cb(userdata, request):
        #traj_request = PlanToTrajectories().Request
        request.pointcloud_in = userdata.segmented_pointcloud_sm
        request.plane_coefficients = userdata.plane_coefficients_sm
        request.target_ellipses = userdata.target_ellipses_sm
        request.actions = userdata.actions_sm
        return request


def main():
        rospy.init_node('smach_example_state_machine')

        # Create a SMACH state machine
        sm = smach.StateMachine(outcomes=['finishOk','failedVision','failedPlanning','failedMoveArm'])
        
        # Open the container
        with sm:
                # Add states to the container
                smach.StateMachine.add('NeedInfo', NeedInfo(),
                                transitions={'next':'GetInfoCamera', 
                                                'failed':'failedVision'},)
                                                
                smach.StateMachine.add('GetInfoCamera', GetInfoCamera(),
                                #ServiceState('/iri_clean_perception/get_board_information/get_board_info_srv',
                                                        #BoardInfo,
                                                        #response_slots=['segmented_image', 
                                                                        #'pointcloud',
                                                                        #'plane_coefficients']),
                                transitions={'succeeded':'LearnerAddState',
                                                'preempted':'failedVision',
                                                'aborted':'failedVision'},
                                remapping={'segmented_image':'segmented_image_sm',
                                                'pointcloud':'segmented_pointcloud_sm',
                                                'plane_coefficients':'plane_coefficients_sm'})
                
                smach.StateMachine.add('LearnerAddState', LearnerAddState(),
                                transitions={'succeeded':'GetStateRepresentation',
                                                'preempted':'failedPlanning',
                                                'aborted':'failedPlanning'},
                                remapping={'rules_updated':'rules_updated_sm'})
                
                
                smach.StateMachine.add('GetStateRepresentation',
                                GetStateRepresentation(),
                                transitions={'is_clean':'finishOk',
                                                'get_full_info':'MoveArmBase',
                                                'succeeded':'PradaPlanning',
                                                'preempted':'failedVision',
                                                'aborted':'failedVision'},
                                remapping={'dirty_areas':'dirty_areas_sm',
                                           'distances':'distances_sm',
                                           'state_string':'state_string_sm'})
                
                smach.StateMachine.add('PradaPlanning',
                                #ServiceState('/iri_clean_plan/prada_planner/getBestAction',
                                                #PradaPlanSrv,
                                                #input_keys=['segmented_pointcloud_sm'],
                                                #output_keys=['segmented_pointcloud_sm'],
                                                #request_slots=['distances','max_widths','min_widths','sparse','num_objects','arm_position'],
                                                #response_slots=['succeeded','target_ellipses','actions']),
                                PradaPlanning(),
                                transitions={'succeeded':'CalculateTrajectories',
                                                'preempted':'failedPlanning',
                                                'aborted':'failedPlanning',
                                                'no_plan':'NeedInfo'},
                                remapping={'target_ellipses':'target_ellipses_sm',
                                           'actions':'actions_sm'})
                
                
                
                smach.StateMachine.add('CalculateTrajectories', 
                                CalculateTrajectories(),
                                #ServiceState('/iri_clean_plan/cleaning_plan/get_trajectories_from_plan',
                                                #PlanToTrajectories,
                                                ##request_slots=['segmented_pointcloud_sm','plane_coefficients_sm','target_ellipses_sm','actions_sm'],
                                                #request_cb = calculate_trajectories_cb,
                                                #input_keys=['segmented_pointcloud_sm','plane_coefficients_sm','target_ellipses_sm','actions_sm']),
                                                ##response_slots=['poses','plane_coefficients','poses_2d_coordinates','secondary_arm']),
                                transitions={'succeeded':'MoveArm',
                                                'preempted':'failedVision',
                                                'aborted':'failedVision',
                                                'no_movements':'NeedInfo'})
                
                
                smach.StateMachine.add('MoveArm', MoveArm(), 
                                transitions={'succeeded':'LearnerAddAction',
                                                'failed':'failedMoveArm'},
                                remapping={'movements_ok':'movements_ok_sm'})
                
                smach.StateMachine.add('LearnerAddAction',
                                LearnerAddAction(),
                                transitions={'needs_learning':'MoveArmBaseAndWait',
                                                'succeeded':'NeedInfo',
                                                'preempted':'failedPlanning',
                                                'aborted':'failedPlanning'})
                
                smach.StateMachine.add('MoveArmBase', MoveArmBase(), 
                                transitions={'next':'NeedInfo', 
                                                'failed':'failedMoveArm'})
                
                smach.StateMachine.add('MoveArmBaseAndWait', MoveArmBaseWait(), 
                                transitions={'next':'WaitASecond', 
                                                'failed':'failedMoveArm'})
                                                
                smach.StateMachine.add('WaitASecond', st_sleep.WaitSeconds(1),
                                transitions={'finish':'NeedInfo'})

        # smach viewer
        sis = smach_ros.IntrospectionServer('iri_board_cleaning', sm, '/SM_LEARNING')
        sis.start()
        
        # Execute SMACH plan
        try:
                outcome = sm.execute()
        finally:
                sis.stop()



if __name__ == '__main__':
        global time_perception
        global time_state
        global time_planning
        global time_learning
        global time_moving
        global time_extra_learn
        
        time_perception = 0
        time_state = 0
        time_planning = 0
        time_learning = 0
        time_moving = 0
        time_extra_learn = 0
        
        start=time.time()
        main()
        finish=time.time()
        total_time = finish-start
        rospy.loginfo('Time taken: ' + str(total_time))
        print "Total time " + str(total_time)
        
        
        
        rospy.loginfo('Time taken perception: ' + str(time_perception))
        rospy.loginfo('Time taken state: ' + str(time_state))
        rospy.loginfo('Time taken planning: ' + str(time_planning))
        rospy.loginfo('Time taken learning: ' + str(time_learning))
        rospy.loginfo('Time taken moving: ' + str(time_moving))
        rospy.loginfo('Time taken extra_learn: ' + str(time_extra_learn))