PythonRobotics/MissionPlanning/StateMachine/robot_behavior_case.py at 8b3ce1058825bd3681242476fb776c136541c93c · Aglargil/PythonRobotics · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
"""
A case study of robot behavior using state machine

author: Wang Zheng (@Aglargil)
"""

from state_machine import StateMachine


class Robot:
    def __init__(self):
        self.battery = 100
        self.task_progress = 0

        # Initialize state machine
        self.machine = StateMachine("robot_sm", self)

        # Add state transition rules
        self.machine.add_transition(
            src_state="patrolling",
            event="detect_task",
            dst_state="executing_task",
            guard=None,
            action=None,
        )

        self.machine.add_transition(
            src_state="executing_task",
            event="task_complete",
            dst_state="patrolling",
            guard=None,
            action="reset_task",
        )

        self.machine.add_transition(
            src_state="executing_task",
            event="low_battery",
            dst_state="returning_to_base",
            guard="is_battery_low",
        )

        self.machine.add_transition(
            src_state="returning_to_base",
            event="reach_base",
            dst_state="charging",
            guard=None,
            action=None,
        )

        self.machine.add_transition(
            src_state="charging",
            event="charge_complete",
            dst_state="patrolling",
            guard=None,
            action="battery_full",
        )

        # Set initial state
        self.machine.set_current_state("patrolling")

    def is_battery_low(self):
        """Battery level check condition"""
        return self.battery < 30

    def reset_task(self):
        """Reset task progress"""
        self.task_progress = 0
        print("[Action] Task progress has been reset")

    # Modify state entry callback naming convention (add state_ prefix)
    def on_enter_executing_task(self):
        print("\n------ Start Executing Task ------")
        print(f"Current battery: {self.battery}%")
        while self.machine.get_current_state().name == "executing_task":
            self.task_progress += 10
            self.battery -= 25
            print(
                f"Task progress: {self.task_progress}%, Remaining battery: {self.battery}%"
            )

            if self.task_progress >= 100:
                self.machine.process("task_complete")
                break
            elif self.is_battery_low():
                self.machine.process("low_battery")
                break

    def on_enter_returning_to_base(self):
        print("\nLow battery, returning to charging station...")
        self.machine.process("reach_base")

    def on_enter_charging(self):
        print("\n------ Charging ------")
        self.battery = 100
        print("Charging complete!")
        self.machine.process("charge_complete")


# Keep the test section structure the same, only modify the trigger method
if __name__ == "__main__":
    robot = Robot()
    print(robot.machine.generate_plantuml())

    print(f"Initial state: {robot.machine.get_current_state().name}")
    print("------------")

    # Trigger task detection event
    robot.machine.process("detect_task")

    print("\n------------")
    print(f"Final state: {robot.machine.get_current_state().name}")