Merge pull request #7 from VirxEC/beta.9

Update to the newest RLBot v5 spec

Author: VirxEC

pyproject.toml

@@ -8,7 +8,7 @@ description = "A high performance Python interface for communicating with RLBot
 dynamic = ["version"]
 requires-python = ">= 3.11"
 dependencies = [
-    "rlbot_flatbuffers~=0.9.1",
+    "rlbot_flatbuffers~=0.10.0",
     "psutil==6.*",
 ]
 readme = "README.md"

rlbot/interface.py

@@ -20,7 +20,7 @@ class SocketDataType(IntEnum):
     """
 
     NONE = 0
-    GAME_TICK_PACKET = 1
+    GAME_PACKET = 1
     FIELD_INFO = 2
     START_COMMAND = 3
     MATCH_SETTINGS = 4
@@ -34,6 +34,7 @@ class SocketDataType(IntEnum):
     STOP_COMMAND = 12
     SET_LOADOUT = 13
     INIT_COMPLETE = 14
+    CONTROLLABLE_TEAM_INFO = 15
 
 
 MAX_SIZE_2_BYTES = 2**16 - 1
@@ -65,16 +66,23 @@ class SocketRelay:
     _should_continue = True
 
     on_connect_handlers: list[Callable[[], None]] = []
-    packet_handlers: list[Callable[[flat.GameTickPacket], None]] = []
+    packet_handlers: list[Callable[[flat.GamePacket], None]] = []
     field_info_handlers: list[Callable[[flat.FieldInfo], None]] = []
     match_settings_handlers: list[Callable[[flat.MatchSettings], None]] = []
     match_communication_handlers: list[Callable[[flat.MatchComm], None]] = []
     ball_prediction_handlers: list[Callable[[flat.BallPrediction], None]] = []
+    controllable_team_info_handlers: list[
+        Callable[[flat.ControllableTeamInfo], None]
+    ] = []
     raw_handlers: list[Callable[[SocketMessage], None]] = []
 
     def __init__(
-        self, connection_timeout: float = 120, logger: Optional[logging.Logger] = None
+        self,
+        agent_id: str,
+        connection_timeout: float = 120,
+        logger: Optional[logging.Logger] = None,
     ):
+        self.agent_id = agent_id
         self.connection_timeout = connection_timeout
         self.logger = get_logger("interface") if logger is None else logger
 
@@ -95,8 +103,8 @@ def send_bytes(self, data: bytes, data_type: SocketDataType):
         message = int_to_bytes(data_type) + int_to_bytes(size) + data
         self.socket.sendall(message)
 
-    def send_init_complete(self, init_complete: flat.InitComplete):
-        self.send_bytes(init_complete.pack(), SocketDataType.INIT_COMPLETE)
+    def send_init_complete(self):
+        self.send_bytes(bytes(), SocketDataType.INIT_COMPLETE)
 
     def send_set_loadout(self, set_loadout: flat.SetLoadout):
         self.send_bytes(set_loadout.pack(), SocketDataType.SET_LOADOUT)
@@ -191,6 +199,7 @@ def connect(
             handler()
 
         flatbuffer = flat.ConnectionSettings(
+            self.agent_id,
             wants_ball_predictions,
             wants_match_communications,
             close_after_match,
@@ -260,9 +269,9 @@ def handle_incoming_message(self, incoming_message: SocketMessage):
         match incoming_message.type:
             case SocketDataType.NONE:
                 self._should_continue = False
-            case SocketDataType.GAME_TICK_PACKET:
+            case SocketDataType.GAME_PACKET:
                 if len(self.packet_handlers) > 0:
-                    packet = flat.GameTickPacket.unpack(incoming_message.data)
+                    packet = flat.GamePacket.unpack(incoming_message.data)
                     for handler in self.packet_handlers:
                         handler(packet)
             case SocketDataType.FIELD_INFO:
@@ -285,6 +294,13 @@ def handle_incoming_message(self, incoming_message: SocketMessage):
                     ball_prediction = flat.BallPrediction.unpack(incoming_message.data)
                     for handler in self.ball_prediction_handlers:
                         handler(ball_prediction)
+            case SocketDataType.CONTROLLABLE_TEAM_INFO:
+                if len(self.controllable_team_info_handlers) > 0:
+                    player_mappings = flat.ControllableTeamInfo.unpack(
+                        incoming_message.data
+                    )
+                    for handler in self.controllable_team_info_handlers:
+                        handler(player_mappings)
 
     def disconnect(self):
         if not self.is_connected:

rlbot/managers/bot.py

@@ -27,20 +27,19 @@ class Bot:
     _initialized_bot = False
     _has_match_settings = False
     _has_field_info = False
+    _has_player_mapping = False
 
-    _latest_packet: Optional[flat.GameTickPacket] = None
+    _latest_packet: Optional[flat.GamePacket] = None
     _latest_prediction = flat.BallPrediction()
 
-    def __init__(self):
-        spawn_id = os.environ.get("RLBOT_SPAWN_IDS")
+    def __init__(self, default_agent_id: Optional[str] = None):
+        agent_id = os.environ.get("RLBOT_AGENT_ID") or default_agent_id
 
-        if spawn_id is None:
-            self.logger.warning("RLBOT_SPAWN_IDS environment variable not set")
-        else:
-            self.spawn_id = int(spawn_id)
-            self.logger.info("Spawn ID: %s", self.spawn_id)
+        if agent_id is None:
+            self.logger.critical("RLBOT_AGENT_ID environment variable is not set")
+            exit(1)
 
-        self._game_interface = SocketRelay(logger=self.logger)
+        self._game_interface = SocketRelay(agent_id, logger=self.logger)
         self._game_interface.match_settings_handlers.append(self._handle_match_settings)
         self._game_interface.field_info_handlers.append(self._handle_field_info)
         self._game_interface.match_communication_handlers.append(
@@ -49,13 +48,23 @@ def __init__(self):
         self._game_interface.ball_prediction_handlers.append(
             self._handle_ball_prediction
         )
+        self._game_interface.controllable_team_info_handlers.append(
+            self._handle_player_mappings
+        )
         self._game_interface.packet_handlers.append(self._handle_packet)
 
         self.renderer = Renderer(self._game_interface)
 
-    def _initialize_agent(self):
+    def _initialize(self):
+        # search match settings for our name
+        for player in self.match_settings.player_configurations:
+            if player.spawn_id == self.spawn_id:
+                self.name = player.name
+                self.logger = get_logger(self.name)
+                break
+
         try:
-            self.initialize_agent()
+            self.initialize()
         except Exception as e:
             self.logger.critical(
                 "Bot %s failed to initialize due the following error: %s", self.name, e
@@ -64,75 +73,53 @@ def _initialize_agent(self):
             exit()
 
         self._initialized_bot = True
-        self._game_interface.send_init_complete(flat.InitComplete(self.spawn_id))
+        self._game_interface.send_init_complete()
 
     def _handle_match_settings(self, match_settings: flat.MatchSettings):
         self.match_settings = match_settings
         self._has_match_settings = True
 
-        # search match settings for our spawn id
-        for player in self.match_settings.player_configurations:
-            if player.spawn_id == self.spawn_id:
-                self.team = player.team
-                self.name = player.name
-                self.logger = get_logger(self.name)
-                break
-
-            if self.spawn_id == 0:
-                match player.variety.item:
-                    case flat.RLBot():
-                        self.team = player.team
-                        self.name = player.name
-                        self.logger = get_logger(self.name)
-                        break
-
-        if not self._initialized_bot and self._has_field_info:
-            self._initialize_agent()
+        if (
+            not self._initialized_bot
+            and self._has_field_info
+            and self._has_player_mapping
+        ):
+            self._initialize()
 
     def _handle_field_info(self, field_info: flat.FieldInfo):
         self.field_info = field_info
         self._has_field_info = True
 
-        if not self._initialized_bot and self._has_match_settings:
-            self._initialize_agent()
+        if (
+            not self._initialized_bot
+            and self._has_match_settings
+            and self._has_player_mapping
+        ):
+            self._initialize()
+
+    def _handle_player_mappings(self, player_mappings: flat.ControllableTeamInfo):
+        self.team = player_mappings.team
+        controllable = player_mappings.controllables[0]
+        self.spawn_id = controllable.spawn_id
+        self.index = controllable.index
+        self._has_player_mapping = True
+
+        if (
+            not self._initialized_bot
+            and self._has_match_settings
+            and self._has_field_info
+        ):
+            self._initialize()
 
     def _handle_ball_prediction(self, ball_prediction: flat.BallPrediction):
         self._latest_prediction = ball_prediction
 
-    def _handle_packet(self, packet: flat.GameTickPacket):
+    def _handle_packet(self, packet: flat.GamePacket):
         self._latest_packet = packet
 
-    def _packet_processor(self, packet: flat.GameTickPacket):
-        if (
-            self.index == -1
-            or len(packet.players) <= self.index
-            or packet.players[self.index].spawn_id != self.spawn_id
-        ):
-            # spawn id should only be 0 if RLBOT_SPAWN_IDS was not set
-            if self.spawn_id == 0:
-                # in this case, if there's only one player, we can assume it's us
-                player_index = -1
-                for i, player in enumerate(packet.players):
-                    # skip human players/psyonix bots
-                    if not player.is_bot:
-                        continue
-
-                    if player_index != -1:
-                        self.logger.error(
-                            "Multiple bots in the game, please set RLBOT_SPAWN_IDS"
-                        )
-                        return
-
-                    player_index = i
-                self.index = player_index
-
-            for i, player in enumerate(packet.players):
-                if player.spawn_id == self.spawn_id:
-                    self.index = i
-                    break
-
-            if self.index == -1:
-                return
+    def _packet_processor(self, packet: flat.GamePacket):
+        if len(packet.players) <= self.index:
+            return
 
         self.ball_prediction = self._latest_prediction
 
@@ -162,9 +149,9 @@ def run(
 
             # custom message handling logic
             # this reads all data in the socket until there's no more immediately available
-            # checks if there was a GameTickPacket in the data, and if so, processes it
+            # checks if there was a GamePacket in the data, and if so, processes it
             # then sets the socket to non-blocking and waits for more data
-            # if there was no GameTickPacket, it sets to blocking and waits for more data
+            # if there was no GamePacket, it sets to blocking and waits for more data
             while True:
                 try:
                     self._game_interface.handle_incoming_messages(True)
@@ -271,23 +258,21 @@ def set_loadout(self, loadout: flat.PlayerLoadout, spawn_id: int):
         """
         Sets the loadout of a bot.
 
-        For use as a loadout generator, call inside of `initialize_agent`.
-        Will be ignored if called outside of `initialize_agent` when state setting is disabled.
+        For use as a loadout generator, call inside of `initialize`.
+        Will be ignored if called outside of `initialize` when state setting is disabled.
         """
         self._game_interface.send_set_loadout(flat.SetLoadout(spawn_id, loadout))
 
-    def initialize_agent(self):
+    def initialize(self):
         """
         Called for all heaver initialization that needs to happen.
-        Field info and match settings are fully loaded at this point, and won't return garbage data.
-
-        NOTE: `self.index` is not set at this point, and should not be used. `self.team` and `self.name` _are_ set with correct information.
+        Field info, match settings, name, index, and team are fully loaded at this point, and won't return garbage data.
         """
 
     def retire(self):
         """Called after the game ends"""
 
-    def get_output(self, packet: flat.GameTickPacket) -> flat.ControllerState:
+    def get_output(self, packet: flat.GamePacket) -> flat.ControllerState:
         """
         Where all the logic of your bot gets its input and returns its output.
         """