Merge pull request #8 from NicEastvillage/script_and_rects

Fix Script manager and add draw_rect_2d and draw_rect_3d

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.10.0",
+    "rlbot_flatbuffers~=0.11.0",
     "psutil==6.*",
 ]
 readme = "README.md"

rlbot/managers/rendering.py

@@ -23,19 +23,19 @@ def _get_anchor(
 
 class Renderer:
     transparent = flat.Color()
-    black = flat.Color(255)
+    black = flat.Color(a=255)
     white = flat.Color(255, 255, 255, 255)
-    grey = gray = flat.Color(255, 128, 128, 128)
-    blue = flat.Color(255, 0, 0, 255)
-    red = flat.Color(255, 255, 0, 0)
-    green = flat.Color(255, 0, 128, 0)
-    lime = flat.Color(255, 0, 255, 0)
-    yellow = flat.Color(255, 255, 255, 0)
-    orange = flat.Color(255, 225, 128, 0)
-    cyan = flat.Color(255, 0, 255, 255)
-    pink = flat.Color(255, 255, 0, 255)
-    purple = flat.Color(255, 128, 0, 128)
-    teal = flat.Color(255, 0, 128, 128)
+    grey = gray = flat.Color(128, 128, 128, 255)
+    blue = flat.Color(0, 0, 255, 255)
+    red = flat.Color(255, 0, 0, 255)
+    green = flat.Color(0, 128, 0, 255)
+    lime = flat.Color(0, 255, 0, 255)
+    yellow = flat.Color(255, 255, 0, 255)
+    orange = flat.Color(225, 128, 0, 255)
+    cyan = flat.Color(0, 255, 255, 255)
+    pink = flat.Color(255, 0, 255, 255)
+    purple = flat.Color(128, 0, 128, 255)
+    teal = flat.Color(0, 128, 128, 255)
 
     _logger = get_logger("renderer")
 
@@ -114,7 +114,7 @@ def is_rendering(self):
         return self._group_id is not None
 
     def _draw(
-        self, render: flat.String2D | flat.String3D | flat.Line3D | flat.PolyLine3D
+        self, render: flat.String2D | flat.String3D | flat.Line3D | flat.PolyLine3D | flat.Rect2D | flat.Rect3D
     ):
         self._current_renders.append(flat.RenderMessage(render))
 
@@ -178,3 +178,39 @@ def draw_string_2d(
                 v_align,
             )
         )
+
+    def draw_rect_2d(
+            self,
+            x: float,
+            y: float,
+            width: float,
+            height: float,
+            color: flat.Color,
+            centered: bool = True
+    ):
+        self._draw(
+            flat.Rect2D(
+                x,
+                y,
+                width,
+                height,
+                color,
+                centered,
+            )
+        )
+
+    def draw_rect_3d(
+            self,
+            anchor: flat.RenderAnchor | flat.BallAnchor | flat.CarAnchor | flat.Vector3,
+            width: float,
+            height: float,
+            color: flat.Color,
+    ):
+        self._draw(
+            flat.Rect3D(
+                _get_anchor(anchor),
+                width,
+                height,
+                color,
+            )
+        )

rlbot/managers/script.py

@@ -18,15 +18,15 @@ class Script:
     index: int = 0
     name: str = "Unknown"
     spawn_id: int = 0
+    agent_id: str = None
 
     match_settings = flat.MatchSettings()
     field_info = flat.FieldInfo()
     ball_prediction = flat.BallPrediction()
 
-    _initialized_bot = False
+    _initialized_script = False
     _has_match_settings = False
     _has_field_info = False
-    _has_player_mapping = False
 
     _latest_packet: Optional[flat.GamePacket] = None
     _latest_prediction = flat.BallPrediction()
@@ -47,15 +47,14 @@ def __init__(self, default_agent_id: Optional[str] = None):
         self._game_interface.ball_prediction_handlers.append(
             self._handle_ball_prediction
         )
-        self._game_interface.controllable_team_info_handlers.append(
-            self._handle_controllable_team_info
-        )
         self._game_interface.packet_handlers.append(self._handle_packet)
 
         self.renderer = Renderer(self._game_interface)
 
-    def _initialize(self):
-        self.name = self.match_settings.script_configurations[self.index].name
+    def _try_initialize(self):
+        if self._initialized_script or not self._has_match_settings:
+            return
+
         self.logger = get_logger(self.name)
 
         try:
@@ -69,46 +68,27 @@ def _initialize(self):
             print_exc()
             exit()
 
-        self._initialized_bot = True
+        self._initialized_script = True
         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
 
-        if (
-            not self._initialized_bot
-            and self._has_field_info
-            and self._has_player_mapping
-        ):
-            self._initialize()
+        for i, script in enumerate(match_settings.script_configurations):
+            if script.agent_id == self.agent_id:
+                self.index = i
+                self.name = script.name
+                self._has_match_settings = True
+                break
+        else:   # else block runs if break was not hit
+            self.logger.warning("Script with agent id '%s' did not find itself in the match settings", self.agent_id)
+
+        self._try_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
-            and self._has_player_mapping
-        ):
-            self._initialize()
-
-    def _handle_controllable_team_info(
-        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()
+        self._try_initialize()
 
     def _handle_ball_prediction(self, ball_prediction: flat.BallPrediction):
         self._latest_prediction = ball_prediction
@@ -117,8 +97,6 @@ def _handle_packet(self, packet: flat.GamePacket):
         self._latest_packet = packet
 
     def _packet_processor(self, packet: flat.GamePacket):
-        if len(packet.players) <= self.index:
-            return
 
         self.ball_prediction = self._latest_prediction
 

tests/render_test/render.py

@@ -1,30 +1,32 @@
 from rlbot import flat
-from rlbot.flat import BallAnchor, Vector3, CarAnchor
+from rlbot.flat import BallAnchor, Vector3, CarAnchor, RenderAnchor, Color
 from rlbot.managers import Script
 
 
 class RenderFun(Script):
     needs_render = True
     last_state = flat.GameStatus.Inactive
+    player_count = 0
 
     def handle_packet(self, packet: flat.GamePacket):
         if (
             packet.game_info.game_status != flat.GameStatus.Replay
             and self.last_state == flat.GameStatus.Replay
-        ):
+        ) or len(packet.players) != self.player_count:
             self.needs_render = True
         self.last_state = packet.game_info.game_status
 
         if self.needs_render:
             self.needs_render = False
+            self.player_count = len(packet.players)
 
-            match packet.balls[0].shape.item:
-                case flat.SphereShape() | flat.CylinderShape() as shape:
-                    radius = shape.diameter / 2
-                case flat.BoxShape() as shape:
-                    radius = shape.length / 2
-                case _:
-                    radius = 0
+            radius = 0
+            if len(packet.balls) > 0:
+                match packet.balls[0].shape.item:
+                    case flat.SphereShape() | flat.CylinderShape() as shape:
+                        radius = shape.diameter / 2
+                    case flat.BoxShape() as shape:
+                        radius = shape.length / 2
 
             self.do_render(radius)
 
@@ -55,6 +57,23 @@ def do_render(self, radius: float):
         for i in range(1, len(points)):
             self.renderer.draw_line_3d(points[i - 1], points[i], self.renderer.red)
 
+        self.renderer.draw_rect_3d(RenderAnchor(Vector3(0, 0, 100), CarAnchor(0, Vector3(200, 0, 0))), 0.02, 0.02, self.renderer.blue)
+        self.renderer.draw_rect_3d(CarAnchor(0, Vector3(200, 0, 0)), 0.02, 0.02, self.renderer.blue)
+
+        self.renderer.draw_rect_2d(0.75, 0.75, 0.1, 0.1, Color(255, 150, 30, 100), centered=False)
+        self.renderer.draw_rect_2d(0.75, 0.75, 0.1, 0.1, self.renderer.black)
+        for hkey, h in {
+            'left': flat.TextHAlign.Left,
+            'center': flat.TextHAlign.Center,
+            'right': flat.TextHAlign.Right,
+        }.items():
+            for vkey, v in {
+                'top': flat.TextVAlign.Top,
+                'center': flat.TextVAlign.Center,
+                'bottom': flat.TextVAlign.Bottom,
+            }.items():
+                self.renderer.draw_string_2d(f'\n\n{vkey:^14}\n{hkey:^14}\n\n', 0.75, 0.75, 0.66, self.renderer.white, h_align=h, v_align=v)
+
         self.renderer.end_rendering()