add state machine

state_machine test

state_machine_update

Author: Aglargil

MissionPlanning/pyrobotics_state_machine.py

@@ -0,0 +1,328 @@
+from collections.abc import Callable
+
+
+class State:
+    def __init__(self, name, on_enter=None, on_exit=None):
+        self.name = name
+        self.on_enter = on_enter
+        self.on_exit = on_exit
+
+    def enter(self):
+        print(f"entering <{self.name}>")
+        if self.on_enter:
+            self.on_enter()
+
+    def exit(self):
+        print(f"exiting <{self.name}>")
+        if self.on_exit:
+            self.on_exit()
+
+
+class StateMachine(State):
+    def __init__(self, model: object, name: str, on_enter=None, on_exit=None):
+        State.__init__(self, name, on_enter, on_exit)
+        self.states = {}
+        self.events = {}
+        self.transition_table = {}
+        self._model = model
+        self._state: StateMachine = None
+
+    def add_transition(
+        self,
+        src_state: str | State,
+        event: str,
+        dst_state: str | State,
+        guard: str | Callable = None,
+        action: str | Callable = None,
+    ) -> None:
+        """Add a transition to the state machine.
+
+        Args:
+            src_state: Source state name or State object
+            event: Event name or Event object
+            dst_state: Destination state name or State object
+            guard: Guard function name or callable
+            action: Action function name or callable
+        """
+        # Convert string parameters to objects if necessary
+        self.register_state(src_state)
+        self.register_event(event)
+        self.register_state(dst_state)
+
+        def get_state_obj(state):
+            return state if isinstance(state, State) else self.get_state(state)
+
+        def get_callable(func):
+            return func if callable(func) else getattr(self._model, func, None)
+
+        src_state_obj = get_state_obj(src_state)
+        dst_state_obj = get_state_obj(dst_state)
+
+        guard_func = get_callable(guard) if guard else None
+        action_func = get_callable(action) if action else None
+        self.transition_table[(src_state_obj.name, event)] = (
+            dst_state_obj,
+            guard_func,
+            action_func,
+        )
+
+    def state_transition(self, src_state: State, event: str):
+        if (src_state.name, event) not in self.transition_table:
+            raise ValueError(
+                f"|{self.name}| invalid transition: <{src_state.name}> : [{event}]"
+            )
+
+        dst_state, guard, action = self.transition_table[(src_state.name, event)]
+
+        def call_guard(guard):
+            if callable(guard):
+                return guard()
+            else:
+                return True
+
+        def call_action(action):
+            if callable(action):
+                action()
+
+        if call_guard(guard):
+            call_action(action)
+            if src_state.name != dst_state.name:
+                print(
+                    f"|{self.name}| transitioning from <{src_state.name}> to <{dst_state.name}>"
+                )
+                src_state.exit()
+                self._state = dst_state
+                dst_state.enter()
+        else:
+            print(
+                f"|{self.name}| skipping transition from <{src_state.name}> to <{dst_state.name}> because guard failed"
+            )
+
+    def register_state(self, state: str | State, on_enter=None, on_exit=None):
+        """Register a state in the state machine.
+
+        Args:
+            state (str | State): The state to register. Can be either a string (state name)
+                                or a State object.
+            on_enter (Callable, optional): Callback function to be executed when entering the state.
+                                         If state is a string and on_enter is None, it will look for
+                                         a method named 'on_enter_<state>' in the model.
+            on_exit (Callable, optional): Callback function to be executed when exiting the state.
+                                        If state is a string and on_exit is None, it will look for
+                                        a method named 'on_exit_<state>' in the model.
+
+        Raises:
+            ValueError: If a state with the same name is already registered with a different type.
+        """
+        if isinstance(state, str):
+            if on_enter is None:
+                on_enter = getattr(self._model, "on_enter_" + state, None)
+            if on_exit is None:
+                on_exit = getattr(self._model, "on_exit_" + state, None)
+            self.states[state] = State(state, on_enter, on_exit)
+            return
+
+        name = state.name
+        if name in self.states and type(self.states[name]) is not type(state):
+            raise ValueError(
+                f'State "{name}" {type(state).__name__} already registered as {type(self.states[name]).__name__}'
+            )
+
+        self.states[name] = state
+
+    def register_event(self, event: str):
+        self.events[event] = event
+
+    def get_state(self, name):
+        return self.states[name]
+
+    def get_event(self, name):
+        return self.events[name]
+
+    def has_event(self, event: str):
+        return event in self.events
+
+    def set_current_state(self, state: State | str):
+        if isinstance(state, str):
+            self._state = self.get_state(state)
+        else:
+            self._state = state
+
+    def get_current_state(self):
+        return self._state
+
+    def process(self, event: str) -> None:
+        """Process an event in the state machine.
+
+        Args:
+            event: Event name or Event object
+        """
+        if self._state is None:
+            raise ValueError("State machine is not initialized")
+
+        if self.has_event(event):
+            self.state_transition(self._state, event)
+        else:
+            raise ValueError(f"Invalid event: {event}")
+
+
+class EventBus:
+    def __init__(self):
+        self.subscribers = {}
+
+    def subscribe(self, event: str, callback: Callable):
+        if event not in self.subscribers:
+            self.subscribers[event] = []
+        self.subscribers[event].append(callback)
+
+    def publish(self, event: str):
+        if event in self.subscribers:
+            for callback in self.subscribers[event]:
+                callback()
+        else:
+            raise ValueError(f"Invalid event: {event}")
+
+
+class SlamModel:
+    def __init__(self, event_bus: EventBus, mapping_success: bool = False):
+        self.event_bus = event_bus
+        self.mapping_success = mapping_success
+
+    def on_enter_localization(self):
+        self.event_bus.publish("top_localization_ready_event")
+
+    def on_enter_mapping(self):
+        self.mapping_success = True
+
+    def is_mapping_success(self):
+        return self.mapping_success
+
+
+class PlanningModel:
+    def __init__(self, event_bus: EventBus):
+        self.event_bus = event_bus
+
+    def on_exit_working(self):
+        self.event_bus.publish("top_stop_working_event")
+
+
+class TopModel:
+    def __init__(self, event_bus: EventBus):
+        self.event_bus = event_bus
+
+    def on_enter_pre_working(self):
+        self.event_bus.publish("slam_start_localization_event")
+
+    def on_enter_working(self):
+        self.event_bus.publish("planning_start_working_event")
+
+    def on_enter_mapping(self):
+        self.event_bus.publish("planning_start_remote_control_control_event")
+        self.event_bus.publish("slam_start_mapping_event")
+
+    def on_exit_mapping(self):
+        self.event_bus.publish("planning_stop_remote_control_control_event")
+        self.event_bus.publish("slam_stop_mapping_event")
+
+
+def main():
+    event_bus = EventBus()
+
+    slam_model = SlamModel(event_bus)
+    planning_model = PlanningModel(event_bus)
+    top_model = TopModel(event_bus)
+
+    slam_machine = StateMachine(slam_model, "slam_machine")
+    planning_machine = StateMachine(planning_model, "planning_machine")
+    top_machine = StateMachine(top_model, "top_machine")
+
+    # fmt: off
+    slam_machine.add_transition("idle", "start_localization_event", "localization", "is_mapping_success")
+    slam_machine.add_transition("localization", "stop_localization_event", "idle")
+    slam_machine.add_transition("idle", "start_mapping_event", "mapping")
+    slam_machine.add_transition("mapping", "stop_mapping_event", "idle")
+
+    planning_machine.add_transition("idle", "start_working_event", "working")
+    planning_machine.add_transition("idle", "stop_working_event", "idle")
+    planning_machine.add_transition("working", "stop_working_event", "idle")
+    planning_machine.add_transition("idle", "start_remote_control_control_event", "remote_control_control")
+    planning_machine.add_transition("remote_control_control", "stop_remote_control_control_event", "idle")
+
+    top_machine.add_transition("idle", "start_working_event", "pre_working")
+    top_machine.add_transition("pre_working", "localization_ready_event", "working")
+    top_machine.add_transition("pre_working", "stop_working_event", "idle")
+    top_machine.add_transition("working", "stop_working_event", "idle")
+    top_machine.add_transition("idle", "start_mapping_event", "mapping")
+    top_machine.add_transition("mapping", "stop_mapping_event", "idle")
+    # fmt: on
+    event_bus.subscribe(
+        "slam_start_localization_event",
+        lambda: slam_machine.process("start_localization_event"),
+    )
+    event_bus.subscribe(
+        "slam_start_mapping_event",
+        lambda: slam_machine.process("start_mapping_event"),
+    )
+    event_bus.subscribe(
+        "slam_stop_mapping_event",
+        lambda: slam_machine.process("stop_mapping_event"),
+    )
+    event_bus.subscribe(
+        "planning_start_working_event",
+        lambda: planning_machine.process("start_working_event"),
+    )
+    event_bus.subscribe(
+        "planning_start_remote_control_control_event",
+        lambda: planning_machine.process("start_remote_control_control_event"),
+    )
+    event_bus.subscribe(
+        "planning_stop_remote_control_control_event",
+        lambda: planning_machine.process("stop_remote_control_control_event"),
+    )
+    event_bus.subscribe(
+        "top_localization_ready_event",
+        lambda: top_machine.process("localization_ready_event"),
+    )
+    event_bus.subscribe(
+        "top_mapping_ready_event",
+        lambda: top_machine.process("mapping_ready_event"),
+    )
+    event_bus.subscribe(
+        "top_stop_working_event",
+        lambda: top_machine.process("stop_working_event"),
+    )
+
+    def working_task():
+        slam_machine.set_current_state("idle")
+        planning_machine.set_current_state("idle")
+        top_machine.set_current_state("idle")
+        # User sends start working event
+        top_machine.process("start_working_event")
+
+        # Planning Model finish the task, and send stop working event
+        planning_machine.process("stop_working_event")
+
+        print("top_machine: ", top_machine.get_current_state().name)
+        print("planning_machine: ", planning_machine.get_current_state().name)
+        print("slam_machine: ", slam_machine.get_current_state().name)
+
+    working_task()
+
+    def mapping_task():
+        slam_machine.set_current_state("idle")
+        planning_machine.set_current_state("idle")
+        top_machine.set_current_state("idle")
+        # User sends start mapping event
+        top_machine.process("start_mapping_event")
+        # User sends stop mapping event
+        top_machine.process("stop_mapping_event")
+
+        print("top_machine: ", top_machine.get_current_state().name)
+        print("planning_machine: ", planning_machine.get_current_state().name)
+        print("slam_machine: ", slam_machine.get_current_state().name)
+
+    mapping_task()
+
+
+if __name__ == "__main__":
+    main()