Make FSMs in gameplay Enum-based

rj_gameplay/rj_gameplay/play/basic_defense.py

@@ -10,6 +10,13 @@
 import stp.role.cost
 from rj_msgs.msg import RobotIntent
 
+from enum import Enum, auto
+
+
+class State(Enum):
+    INIT = auto()
+    ACTIVE = auto()
+
 
 class BasicDefense(stp.play.Play):
     """Play that consists of:
@@ -21,23 +28,21 @@ class BasicDefense(stp.play.Play):
     def __init__(self):
         super().__init__()
 
-        # super simple FSM
-        # TODO: use FSM class (or at least don't use string literals)
-        self._state = "init"
+        self._state = State.INIT
 
     def tick(
         self,
         world_state: stp.rc.WorldState,
     ) -> List[RobotIntent]:
 
-        if self._state == "init":
+        if self._state == State.INIT:
             self.prioritized_tactics.append(goalie_tactic.GoalieTactic(world_state, 0))
             self.prioritized_tactics.append(wall_tactic.WallTactic(world_state, 5))
             # TODO: add nmark tactic
             #       and make it go for the ball (rather than stopping in front)
             self.assign_roles(world_state)
-            self._state = "active"
+            self._state = State.ACTIVE
             return self.get_robot_intents(world_state)
-        elif self._state == "active":
+        elif self._state == State.ACTIVE:
             # return robot intents from assigned tactics back to gameplay node
             return self.get_robot_intents(world_state)

rj_gameplay/rj_gameplay/play/keepaway.py

@@ -6,6 +6,14 @@
 
 from rj_msgs.msg import RobotIntent
 
+from enum import Enum, auto
+
+
+class State(Enum):
+    INIT = auto()
+    ACTIVE = auto()
+    ASSIGN_ROLES = auto()
+
 
 class Keepaway(stp.play.Play):
     """Play that passes repeatedly, effectively playing keepaway.
@@ -15,9 +23,7 @@ class Keepaway(stp.play.Play):
     def __init__(self):
         super().__init__()
 
-        # super simple FSM
-        # TODO: use FSM class (or at least don't use string literals)
-        self._state = "init"
+        self._state = State.INIT
 
     def tick(
         self,
@@ -31,33 +37,33 @@ def tick(
         (the effect is to pass indefinitely)
         """
 
-        if self._state == "init":
+        if self._state == State.INIT:
             self.prioritized_tactics = [pass_tactic.PassTactic(world_state)]
             # TODO: either add seek tactic(s) or unassigned behavior
 
             self.assign_roles(world_state)
-            self._state = "active"
+            self._state = State.ACTIVE
             return self.get_robot_intents(world_state)
 
-        elif self._state == "active":
+        elif self._state == State.ACTIVE:
             # TODO: this loop's logic is fairly crucial in role assignment
             #
             # is there a way I can force this to happen as a precondition to assign_roles?
             # maybe call assign_roles() every tick but check tactic for needs_assign before assigning it
             # (this works as the method is in Play superclass)
             for tactic in self.prioritized_tactics:
                 if tactic.needs_assign:
-                    self._state = "assign_roles"
+                    self._state = State.ASSIGN_ROLES
 
             # only one tactic in this play
             tactic = self.prioritized_tactics[0]
             if tactic.is_done(world_state):
-                self._state = "init"
+                self._state = State.INIT
 
             return self.get_robot_intents(world_state)
 
-        elif self._state == "assign_roles":
+        elif self._state == State.ASSIGN_ROLES:
             # duplicate code from init
             self.assign_roles(world_state)
-            self._state = "active"
+            self._state = State.ACTIVE
             return self.get_robot_intents(world_state)

rj_gameplay/rj_gameplay/play/line_up.py

@@ -17,31 +17,37 @@
 import numpy as np
 from rj_msgs.msg import RobotIntent
 
+from enum import Enum, auto
+
+
+class State(Enum):
+    INIT = auto()
+    LINE_UP = auto()
+    DONE = auto()
+
 
 class LineUp(stp.play.Play):
     """Lines up all six robots on the side of the field."""
 
     def __init__(self):
         super().__init__()
 
-        # super simple FSM
-        # TODO: use FSM class (or at least don't use string literals)
-        self._state = "init"
+        self._state = State.INIT
 
     def tick(
         self,
         world_state: stp.rc.WorldState,
     ) -> List[RobotIntent]:
 
-        if self._state == "init":
+        if self._state == State.INIT:
             self.prioritized_tactics.append(line_tactic.LineTactic(world_state))
             self.assign_roles(world_state)
-            self._state = "line_up"
+            self._state = State.LINE_UP
             return self.get_robot_intents(world_state)
-        elif self._state == "line_up":
+        elif self._state == State.LINE_UP:
             if self.prioritized_tactics[0].is_done(world_state):
-                self._state = "done"
+                self._state = State.DONE
             return self.get_robot_intents(world_state)
-        elif self._state == "done":
+        elif self._state == State.DONE:
             # TODO: does this state need to exist?
             return None

rj_gameplay/rj_gameplay/role/passer.py

@@ -5,6 +5,17 @@
 
 from rj_msgs.msg import RobotIntent
 
+from enum import Enum, auto
+
+
+class State(Enum):
+    INIT = auto()
+    CAPTURING = auto()
+    PASS_READY = auto()
+    INIT_EXECUTE_PASS = auto()
+    EXECUTE_PASS = auto()
+    KICK_DONE = auto()
+
 
 class PasserRole(stp.role.Role):
     def __init__(self, robot: stp.rc.Robot) -> None:
@@ -13,17 +24,16 @@ def __init__(self, robot: stp.rc.Robot) -> None:
         self.receive_skill = None
         self.pivot_kick_skill = None
 
-        # TODO: make FSM class (or at least use enum instead of str literals)
-        self._state = "init"
+        self._state = State.INIT
 
         self._target_point = None
 
     @property
     def pass_ready(self):
-        return self._state == "pass_ready"
+        return self._state == State.PASS_READY
 
     def set_execute_pass(self, target_point):
-        self._state = "init_execute_pass"
+        self._state = State.INIT_EXECUTE_PASS
         self._target_point = target_point
 
     def tick(self, world_state: stp.rc.WorldState) -> RobotIntent:
@@ -35,35 +45,35 @@ def tick(self, world_state: stp.rc.WorldState) -> RobotIntent:
         """
 
         intent = None
-        if self._state == "init":
+        if self._state == State.INIT:
             self.receive_skill = receive.Receive(robot=self.robot)
             intent = self.receive_skill.tick(world_state)
-            self._state = "capturing"
-        elif self._state == "capturing":
+            self._state = State.CAPTURING
+        elif self._state == State.CAPTURING:
             intent = self.receive_skill.tick(world_state)
             if self.receive_skill.is_done(world_state):
-                self._state = "pass_ready"
-        elif self._state == "pass_ready":
+                self._state = State.PASS_READY
+        elif self._state == State.PASS_READY:
             # TODO: dribble until the receiver is ready
             pass
         # this state transition is done by the PassTactic, which is not canonical FSM
-        elif self._state == "init_execute_pass":
+        elif self._state == State.INIT_EXECUTE_PASS:
             # TODO: make these params configurable
             self.pivot_kick_skill = pivot_kick.PivotKick(
                 robot=self.robot,
                 target_point=self._target_point,
                 chip=False,
                 kick_speed=4.0,  # TODO: adjust based on dist from target_point
             )
-            self._state = "execute_pass"
-        elif self._state == "execute_pass":
+            self._state = State.EXECUTE_PASS
+        elif self._state == State.EXECUTE_PASS:
             intent = self.pivot_kick_skill.tick(world_state)
 
             if self.pivot_kick_skill.is_done(world_state):
-                self._state = "kick_done"
+                self._state = State.KICK_DONE
                 # end FSM
 
         return intent
 
     def is_done(self, world_state) -> bool:
-        return self._state == "kick_done"
+        return self._state == State.KICK_DONE

rj_gameplay/rj_gameplay/role/receiver.py

@@ -5,24 +5,32 @@
 
 from rj_msgs.msg import RobotIntent
 
+from enum import Enum, auto
+
+
+class State(Enum):
+    INIT = auto()
+    PASS_READY = auto()
+    RECEIVE_PASS = auto()
+    DONE = auto()
+
 
 class ReceiverRole(stp.role.Role):
     def __init__(self, robot: stp.rc.Robot) -> None:
         super().__init__(robot)
 
         self.receive_skill = None
 
-        # TODO: make FSM class (or at least use enum instead of str literals)
-        self._state = "init"
+        self._state = State.INIT
 
         self._target_point = None
 
     @property
     def pass_ready(self):
-        return self._state == "pass_ready"
+        return self._state == State.PASS_READY
 
     def set_receive_pass(self):
-        self._state = "receive_pass"
+        self._state = State.RECEIVE_PASS
         self.receive_skill = receive.Receive(robot=self.robot)
 
     def tick(self, world_state: stp.rc.WorldState) -> RobotIntent:
@@ -35,16 +43,16 @@ def tick(self, world_state: stp.rc.WorldState) -> RobotIntent:
 
         intent = None
 
-        if self._state == "init":
-            # do seek behavior
+        if self._state == State.INIT:
+            # TODO: do seek behavior
             pass
-        elif self._state == "receive_pass":
+        elif self._state == State.RECEIVE_PASS:
             intent = self.receive_skill.tick(world_state)
             if self.receive_skill.is_done(world_state):
-                self._state = "done"
+                self._state = State.DONE
                 # end FSM
 
         return intent
 
     def is_done(self, world_state) -> bool:
-        return self._state == "done"
+        return self._state == State.DONE

rj_gameplay/rj_gameplay/skill/pivot_kick.py

@@ -14,6 +14,15 @@
 import numpy as np
 from rj_gameplay.MAX_KICK_SPEED import MAX_KICK_SPEED
 
+from enum import Enum, auto
+
+
+class State(Enum):
+    CAPTURE = auto()
+    PIVOT = auto()
+    KICK = auto()
+    DONE = auto()
+
 
 class PivotKick(skill.Skill):  # add ABC if fails
     """
@@ -55,30 +64,29 @@ def __init__(
         )
         self.capture = capture.Capture(robot)
 
-        self._state = "capture"
+        self._state = State.CAPTURE
 
     def tick(self, world_state: rc.WorldState) -> RobotIntent:
         super().tick(world_state)
-        print("pivot kick state:", self._state)
 
         intent = None
-        if self._state == "capture":
+        if self._state == State.CAPTURE:
             intent = self.capture.tick(world_state)
             if self.capture.is_done(world_state):
-                self._state = "pivot"
-        elif self._state == "pivot":
+                self._state = State.PIVOT
+        elif self._state == State.PIVOT:
             intent = self.pivot.tick(world_state)
             if self.pivot.is_done(world_state):
-                self._state = "kick"
-        elif self._state == "kick":
+                self._state = State.KICK
+        elif self._state == State.KICK:
             intent = self.kick.tick(world_state)
             if self.kick.is_done(world_state):
-                self._state = "done"
+                self._state = State.DONE
 
         return intent
 
     def is_done(self, world_state: rc.WorldState) -> bool:
-        return self._state == "done"
+        return self._state == State.DONE
 
     def __str__(self):
         return f"Pivot(robot={self.robot.id if self.robot is not None else '??'}, target={self.target_point})"