single_step_full_bt.py

import time
import logging
from pathlib import Path
import py_trees
import py_trees.console as console
import py_trees.display as display
from py_trees.visitors import SnapshotVisitor

import config
from packages.core.model_trainer import ModelTrainer
from packages.core.data_manager import DataManager
from tools.classifier_evaluation.visualization.learning_curve import LearningCurveVisualizer

from nodes.single_step.wait_for_enter import WaitForEnter
from nodes.single_step.detect_object_node import DetectObjectNode
from nodes.single_step.classify_and_get_action import ClassifyAndGetAction
from nodes.single_step.execute_action import ExecuteAction
from nodes.single_step.handle_feedback import HandleFeedback

from packages.core.keyboard_listener import setup_keyboard_listener

from tools.classifier_evaluation.core.feedback_manager import FeedbackManager
from packages.logging.logger import setup_logging
from packages.efficiency.measure_time_cpu import measure_time

# Only import ROS-related modules if using ROS
if config.WITH_ROS:
    import rclpy
    from packages.robot.robot_client import RobotActionClient
else:
    from packages.utils.mock_robot_client import MockRobotActionClient

logger = logging.getLogger("root")
setup_logging()

class BTConfigurator:
    def __init__(self):
        """Initialize the Behavior Tree configurator"""
        # Setup logging
        self.logger = logging.getLogger("root")
        
        # Initialize components as None
        self.data_manager = None
        self.validate_data_manager = None
        self.validation_objects = None
        self.feedback_manager = None
        self.keyboard_state = None
        self.keyboard_thread = None
        self.node = None
        self.robot_client = None
        self.learner = None
        
        # BT components
        self.tree = None
        self.root = None
        self.snapshot_visitor = None
        
        # BT nodes
        self.wait_enter = None
        self.detect_object_node = None
        self.classify_action = None
        self.execute_action = None
        self.wait_feedback = None
        
        # BT blackboard
        self.blackboard = py_trees.blackboard.Client(name="Main BT Blackboard")
        
        # Evaluation metrics
        self.demo_counts = []
        self.accuracies = []
        self.images_dir = config.TASK_IMAGES_OUTPUT
        
        if self.images_dir:
            self.images_dir.mkdir(exist_ok=True)
    
    @measure_time("Setup components in Single-step Full BT: ", config.MEASURE_INTERNAL_TIME)
    def setup_components(self):
        """Set up all required components for the behavior tree"""
        # Load training and validation data
        self.data_manager = DataManager(
            demos_path=config.INITIAL_DEMOS_PATH,
            objects_path=config.OBJECTS_YAML_PATH,
            actions_path=config.ACTIONS_YAML_PATH,
            for_evaluation=False
        )
        self.data_manager.load_data()
        
        self.validate_data_manager = DataManager(
            validation_objects_path=config.VALIDATION_OBJECTS_YAML_PATH,
            for_evaluation=False
        )
        self.validation_objects = self.validate_data_manager.load_validation_data_only()
        
        if self.validation_objects is None:
            self.logger.error("No validation objects found")
            return False
        
        # Initialize feedback manager
        self.feedback_manager = FeedbackManager(Path(""))
        
        # Setup keyboard listener
        self.keyboard_state, self.keyboard_thread = setup_keyboard_listener()

        # Initialize ROS components or mock components
        if config.WITH_ROS:
            rclpy.init()
            self.node = rclpy.create_node("robot_behaviour_tree")
            self.robot_client = RobotActionClient(self.node)
            self.robot_client.start()
        else:
            self.node = None
            self.robot_client = MockRobotActionClient()

        # Initialize model
        self.learner = ModelTrainer(model_type=config.MAIN_MODEL_TYPE)
        self.logger.info(f"Using model type: {config.MAIN_MODEL_TYPE}")
        
        return True
    
    @measure_time("Initial training in Single-step Full BT: ", config.MEASURE_INTERNAL_TIME)
    def initial_training(self):
        """Perform initial model training"""
        self.learner.train(self.data_manager.train_demos)
        self.logger.info(f"Initial training completed with {len(self.data_manager.train_demos)} demonstrations")
        
    @measure_time("Initial validation in Single-step Full BT: ", config.MEASURE_INTERNAL_TIME)
    def initial_validation(self):
        """Perform initial model validation"""
        val_accuracy = self.evaluate_model_on_validation()
        self.demo_counts.append(len(self.data_manager.train_demos))
        self.accuracies.append(val_accuracy)
        self.draw_learning_curve()
        
    @measure_time("Setup tree in Single-step Full BT: ", config.MEASURE_INTERNAL_TIME)
    def setup_tree(self):
        """Create and set up the behavior tree"""
        # Create behavior tree nodes
        self.wait_enter = WaitForEnter(
            config.WAIT_FOR_ENTER_NODE_NAME, 
            self.keyboard_state
            )
        
        self.detect_object_node = DetectObjectNode(
            config.DETECT_OBJECT_NODE_NAME,
            config.MAIN_BT_BLACKBOARD_NAME,
            self.data_manager
        )
        self.detect_object_node.setup(ros_node=self.node)
            
        self.classify_action = ClassifyAndGetAction(
            config.CLASSIFY_AND_GET_ACTION_NODE_NAME,
            config.MAIN_BT_BLACKBOARD_NAME,
            self.learner,
            config.OBJECTS_YAML_PATH,
            config.ACTIONS_YAML_PATH,
            self.data_manager
        )
        
        self.execute_action = ExecuteAction(
            config.EXECUTE_ACTION_NODE_NAME,
            config.MAIN_BT_BLACKBOARD_NAME,
            self.robot_client
        )
        
        self.wait_feedback = HandleFeedback(
            config.WAIT_FEEDBACK_NODE_NAME, 
            config.MAIN_BT_BLACKBOARD_NAME,
            self.keyboard_state, 
            self.data_manager,
            self.feedback_manager,
            self.learner
        )

        # Build a sequence tree
        self.root = py_trees.composites.Sequence("Main Sequence", memory=True)
        self.root.add_children([
            self.wait_enter, 
            self.detect_object_node, 
            self.classify_action, 
            self.execute_action, 
            self.wait_feedback
        ])

        # Create a behavior tree
        self.tree = py_trees.trees.BehaviourTree(self.root)
        
        # Add visitors for logging
        self.snapshot_visitor = SnapshotVisitor()
        self.tree.visitors.append(self.snapshot_visitor)
        
        # Enable py_trees logging
        py_trees.logging.level = py_trees.logging.Level.DEBUG
        
        # Log the initial tree structure
        self.logger.info("\nBehavior Tree Structure:")
        self.logger.info(display.ascii_tree(self.root, visited=self.snapshot_visitor.visited, show_status=True))
    
    def run(self):
        """Run the behavior tree main loop"""
        self.logger.info("*** Starting Behavior Tree ***")
        self.logger.info("Press 'q' at any time to exit the program")
        
        # Main loop
        running = True
        tick_count = 0
        
        while running and not self.keyboard_state.exit_requested:
            # Handle ROS if needed
            if config.WITH_ROS and self.node:
                rclpy.spin_once(self.node, timeout_sec=0.01)
            
            # Log pre-tick state
            self.logger.debug(f"Tick {tick_count}: Tree tick starting")
            
            # Tick the tree
            self.tree.tick()
            tick_count += 1
            
            self.log_tree_structure_and_status(tick_count)

            # If tree completed successfully, test model and draw the accuracy curve
            if self.tree.root.status == py_trees.common.Status.SUCCESS:
                self.logger.info("Tree completed successfully, evaluating model")
                self.evaluate_and_update_curve()
                
                # Reset keyboard states after completion
                self.keyboard_state.enter_pressed = False
                self.keyboard_state.space_pressed = False
            
            time.sleep(0.1)
        
        # Clean up
        self.cleanup()
    
    def log_tree_structure_and_status(self, tick_count):
        """Log the tree structure and status"""
        # Log post-tick state
        self.logger.info(f"Tree state after tick {tick_count}:")
        self.logger.info(display.ascii_tree(self.root, visited=self.snapshot_visitor.visited, show_status=True))
        
        # Log tree status
        status_color = {
            py_trees.common.Status.SUCCESS: console.green,
            py_trees.common.Status.RUNNING: console.blue,
            py_trees.common.Status.FAILURE: console.red
        }.get(self.tree.root.status, console.white)
        
        self.logger.info(f"Tree status: {status_color}")
        
        # Log important node statuses
        for node in [self.wait_enter, self.detect_object_node, self.classify_action, 
                        self.execute_action, self.wait_feedback]:
            if node:  # Make sure the node exists
                self.logger.info(f"Node {node.name} status: {node.status}")
    
    @measure_time("Evaluate model on validation in Single-step Full BT: ", config.MEASURE_INTERNAL_TIME)
    def evaluate_model_on_validation(self):
        """Evaluate the current model on the validation set."""
        if not self.learner.is_fitted:
            self.logger.warning("Model not fitted. Returning 0 accuracy.")
            return 0.0
        
        # Make predictions on validation objects
        predictions, _ = self.learner.predict(self.validation_objects)
        
        # Calculate accuracy
        correct = 0
        total = 0
        
        for obj, pred in zip(self.validation_objects, predictions):
            if obj.target_position is not None:
                total += 1
                if obj.target_position == pred:
                    correct += 1
        
        accuracy = correct / total if total > 0 else 0.0
        self.logger.info(f"Validation accuracy: {accuracy:.4f} ({correct}/{total})")
        
        return accuracy
    
    @measure_time("Evaluate and update curve in Single-step Full BT: ", config.MEASURE_INTERNAL_TIME)
    def evaluate_and_update_curve(self):
        """Evaluate the model and update the learning curve"""
        # Get validation accuracy
        val_accuracy = self.evaluate_model_on_validation()
        
        # Add current data point
        self.demo_counts.append(len(self.data_manager.train_demos))
        self.accuracies.append(val_accuracy)
        
        # Draw and display learning curve
        self.draw_learning_curve()
    
    @measure_time("Draw learning curve in Single-step Full BT: ", config.MEASURE_INTERNAL_TIME)
    def draw_learning_curve(self):
        """Create, save, and display the learning curve."""
        if not self.demo_counts or not self.accuracies:
            self.logger.warning("No data to plot learning curve")
            return
        
        # Create output path
        if not self.images_dir:
            self.logger.warning("No images directory set. Learning curve not saved.")
            return
        
        output_path = self.images_dir / f"robot_learning_curve_{config.MAIN_MODEL_TYPE}.png"
        
        # Create the learning curve
        LearningCurveVisualizer.create_learning_curve(
            demo_counts=self.demo_counts,
            accuracies=self.accuracies,
            output_path=output_path,
            title="Learning Curve - Validation Accuracy",
            xlabel="Number of Training Demonstrations",
            ylabel="Accuracy"
        )
        
        self.logger.debug(f"Saved learning curve to {output_path}")
    
    def cleanup(self):
        """Clean up resources and shut down services"""
        if config.WITH_ROS:
            rclpy.shutdown()

def main():
    # Create and setup the behavior tree configurator
    bt_config = BTConfigurator()
    
    # Setup all components
    if not bt_config.setup_components():
        logger.error("Failed to set up components")
        return
    
    # Initial training
    bt_config.initial_training()
    
    # Initial validation
    bt_config.initial_validation()
    
    # Setup behavior tree
    bt_config.setup_tree()
    
    # Run the behavior tree
    bt_config.run()
    
    # Clean up if using ROS
    bt_config.cleanup()

if __name__ == '__main__':
    main()