Commit c76c38d5 authored by Ardillen66's avatar Ardillen66
Browse files

Created main and baxter agent

parent e53993e8
#Learning imports
import features_detection as feat
from decision_api.decision_tree_gr import DecisionTreeGR
from planning.baxter import Baxter
from planning.scene import Scene
#Baxter and ROS imports
import baxter_interface
import rospy
from baxter_core_msgs.msgs import DigitalIOState, ITBState
from baxter_examples import xdisplay_image
class BaxterAgent(object):
"""
Agent that manages the baxter robot learning process and deployment
"""
def __init__(self, config):
self.config = config
self.baxter = Baxter();
self.scene = Scene(self.baxter.get_planning_frame())
rospy.sleep(2)
self.init_scene()
self.is_train = config['is_train']
if self.is_train:
self.tree = DecisionTreeGR()
self.train()
else :
training_samples = self.load_trainig_samples()
self.tree = DecisionTreeGR(training_samples)
self.run()
def train(self):
self.drawTrainingUI()
while self.is_train:
######
#TODO# take snapshot and extract features if object was added to the scene. Update scene accordingly
######
features = {}
right_button = rospy.wait_for_message('/robot/digital_io/right_shoulder_button/state', DigitalIOState, 3)
left_button = rospy.wait_for_message('/robot/digital_io/left_shoulder_button/state', DigitalIOState, 3)
#Use navigator cancel on left or right arm to stop training
left_nav = rospy.wait_for_message('/robot/navigators/left_itb/state', ITBState, 3)
right_nav = rospy.wait_for_message('/robot/navigators/right_itb/state', ITBState, 3)
if right_button.state == 1:
features['direction'] = 'right'
self.sort_right()
self.tree.add_training_sample(features)
elif left_button.state == 1:
features['direction'] = 'left'
self.sort_left()
self.tree.add_training_sample(features)
elif left_nav.buttons[1] or right_nav.buttons[1]:
self.is_train = False
self.init_move() #Go back to default position
self.save_training_samples() #Store trained samples to retrain decision tree
def run(self):
self.drawProductionUI()
#For now we retrain decision tree from saved training data, but maybe better to save trained tree?
self.tree.train()
run = True
while run:
#TODO# Check if object on table and retrieve features
features = {}
#Listen to navigators for both sides
left_nav = rospy.wait_for_message('/robot/navigators/left_itb/state', ITBState, 3)
right_nav = rospy.wait_for_message('/robot/navigators/right_itb/state', ITBState, 3)
if left_nav.buttons[0] or right_nav.buttons[0]:
side = self.tree.classify(features) #Choose sorting side
if side == 'RIGHT':
self.sort_right()
elif side == 'LEFT':
self.sort_left()
elif left_nav.buttons[1] or right_nav.buttons[1]:
run = False
self.init_move()
def load_trainig_samples(self):
#TODO#: load training samles from file specified in config
pass
def save_training_samples(self):
#TODO#: save decision tree training samples to file specified in config
pass
def init_scene(self):
#TODO#: initialize scene with description in config
self.scene.clear_scene()
self.scene.set_table()
self.scene.set_item()
rospy.sleep(2)
def init_move(self):
self.baxter.both.move(self.baxter.both.neutral)
self.baxter.right.gripper.close()
rospy.sleep(3)
def drawTrainingUI(self):
#TODO#: draw training interface on baxter face screen
def drawProductionUI(self):
#TODO#
def sort_right(self):
self.baxter.right.move(self.baxter.right.pick)
self.baxter.right.gripper.open()
rospy.sleep(3)
self.baxter.right.move(self.baxter.right.grasp)
rospy.sleep(1)
self.baxter.right.gripper.close()
rospy.sleep(1)
self.baxter.right.move(self.baxter.right.pick)
self.baxter.right.move(self.baxter.right.neutral)
self.baxter.right.gripper.open()
def sort_left(self):
self.baxter.left.move(self.baxter.left.pick)
self.baxter.left.gripper.open()
rospy.sleep(3)
self.baxter.left.move(self.baxter.left.grasp)
rospy.sleep(1)
self.baxter.left.gripper.close()
rospy.sleep(1)
self.baxter.left.move(self.baxter.left.pick)
self.baxter.left.move(self.baxter.left.neutral)
self.baxter.left.gripper.open()
import moveit_commander
import rospy
import sys
from baxter_agent import BaxterAgent
config = {
'is_train': True,
'tree_dir': './saved_trees',
'scene': 'scene description' #Should contain a description of the scene to be loaded
}
def main():
moveit_commander.roscpp_initialize(sys.argv)
rospy.init_node('garbage_bot', anonymous=True)
BaxterAgent(config)
if __name__ == '__main__':
main()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment