Refactor Stats Tracking

src/software/evaluation/loggers/BUILD

@@ -0,0 +1,9 @@
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "stats_logger",
+    srcs = ["stats_logger.py"],
+    deps = [
+        "//software/evaluation/trackers:tracker",
+    ],
+)

src/software/evaluation/loggers/stats_logger.py

@@ -0,0 +1,170 @@
+import os
+
+from software.evaluation.trackers import (
+    PossessionTracker,
+    ShotTracker,
+    PassTracker,
+    TrackerBuilder,
+    RefereeTracker,
+    GoalieTracker,
+)
+from dataclasses import dataclass
+from software.thunderscope.proto_unix_io import ProtoUnixIO
+from software.thunderscope.constants import RuntimeManagerConstants
+from software.evaluation.logs.event_log import Team as TeamEnum, EventLog
+import logging
+from proto.import_all_protos import *
+import queue
+
+
+@dataclass
+class FSStats:
+    """Stats for how well a FullSystem is performing"""
+
+    num_yellow_cards: int = 0
+    num_red_cards: int = 0
+    num_scores: int = 0
+
+    num_shots_on_net: int = 0
+    num_enemy_shots_blocked: int = 0
+
+
+class StatsLogger:
+    # From GoalieTacticConfig
+    INCOMING_SHOT_MIN_VELOCITY = 0.2
+
+    EVENT_BUFFER_SIZE = 100
+
+    def __init__(
+        self,
+        proto_unix_io: ProtoUnixIO,
+        friendly_colour_yellow: bool,
+        out_file_name: str | None = None,
+        buffer_size: int = 5,
+        record_enemy_stats: bool = False,
+    ):
+        """Initializes the FullSystem Stats Tracker
+
+        :param friendly_colour_yellow: if the friendly colour is yellow
+        :param out_file_name: name of file to write stats to.
+                              If None, uses the value from constants
+        :param buffer_size: the buffer size for protocol buffers
+        :param record_enemy_stats: if this should record both friendly and enemy stats or just friendly
+        """
+        self.friendly_colour_yellow = friendly_colour_yellow
+
+        self.events_file_path = os.path.join(
+            RuntimeManagerConstants.RUNTIME_EVENTS_DIRECTORY_PATH,
+            RuntimeManagerConstants.RUNTIME_EVENTS_FILE
+            if out_file_name is None
+            else out_file_name,
+        )
+        # initialized in setup()
+        self.events_file_handle = None
+
+        self.event_queue = queue.Queue(self.EVENT_BUFFER_SIZE)
+
+        # flag to turn off logging stats if needed
+        self.logging_enabled = True
+
+        self.tracker = (
+            TrackerBuilder(
+                proto_unix_io=proto_unix_io,
+                from_team=(
+                    TeamEnum.YELLOW if self.friendly_colour_yellow else TeamEnum.BLUE
+                ),
+                event_queue=self.event_queue,
+                buffer_size=buffer_size,
+            )
+            .add_tracker(PassTracker)
+            .add_tracker(ShotTracker)
+            .add_tracker(PossessionTracker)
+            .add_tracker(
+                RefereeTracker,
+                friendly_color_yellow=self.friendly_colour_yellow,
+                toggle_logging=self._toggle_logging,
+            )
+            .add_tracker(GoalieTracker, for_friendly=True)
+        )
+
+        self.record_enemy_stats = record_enemy_stats
+        if self.record_enemy_stats:
+            self.enemy_tracker = (
+                TrackerBuilder(
+                    proto_unix_io=proto_unix_io,
+                    from_team=(
+                        TeamEnum.YELLOW
+                        if self.friendly_colour_yellow
+                        else TeamEnum.BLUE
+                    ),
+                    for_team=(
+                        TeamEnum.BLUE
+                        if self.friendly_colour_yellow
+                        else TeamEnum.YELLOW
+                    ),
+                    event_queue=self.event_queue,
+                    buffer_size=buffer_size,
+                )
+                .add_tracker(
+                    RefereeTracker,
+                    friendly_color_yellow=(not self.friendly_colour_yellow),
+                    toggle_logging=self._toggle_logging,
+                )
+                .add_tracker(GoalieTracker, for_friendly=False)
+            )
+
+    def refresh(self) -> None:
+        """Refreshes the events for the game so far"""
+        self.tracker.refresh()
+
+        if not self.events_file_handle:
+            return
+
+        while not self.event_queue.empty():
+            try:
+                # Get item without blocking
+                event = self.event_queue.get_nowait()
+
+                self._write_event_to_file(event)
+            except queue.Empty:
+                return
+
+    def __enter__(self):
+        """Sets up the file resources for logging
+        Creates any missing directories and stores the file handle
+        """
+        # create temp stats directory if it doesn't exist
+        os.makedirs(os.path.dirname(self.events_file_path), exist_ok=True)
+
+        self.events_file_handle = open(self.events_file_path, "a")
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Writes all logs back to file, and cleans up any created file resources after logging"""
+        if self.events_file_handle:
+            self.events_file_handle.flush()
+            self.events_file_handle.close()
+
+    def _toggle_logging(self, should_log: bool) -> None:
+        """Turns logging off or on based on the given boolean
+
+        ;param should_log: True if logging should continue, False if not
+        """
+        self.logging_enabled = should_log
+
+    def _write_event_to_file(self, event: EventLog) -> None:
+        """Write the given stats to the given file
+
+        :param event: the event to write
+        """
+        if not self.events_file_handle:
+            return
+
+        try:
+            csv_row = event.to_csv_row()
+            self.events_file_handle.write(csv_row + "\n")
+            self.events_file_handle.flush()
+
+        except (IOError, FileNotFoundError, PermissionError):
+            logging.warning("Failed to write event to file")

src/software/evaluation/logs/BUILD

@@ -0,0 +1,39 @@
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "log_interface",
+    srcs = ["log_interface.py"],
+    deps = [
+        "//proto:import_all_protos",
+        "//software/thunderscope:time_provider",
+    ],
+)
+
+py_library(
+    name = "world_state_log",
+    srcs = ["world_state_log.py"],
+    data = [
+        "//software:py_constants.so",
+    ],
+    deps = [
+        ":log_interface",
+    ],
+)
+
+py_library(
+    name = "event_log",
+    srcs = ["event_log.py"],
+    deps = [
+        ":log_interface",
+        ":world_state_log",
+    ],
+)
+
+py_library(
+    name = "logs",
+    deps = [
+        ":event_log",
+        ":log_interface",
+        ":world_state_log",
+    ],
+)

src/software/evaluation/logs/event_log.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from enum import StrEnum, auto
+from proto.import_all_protos import *
+from typing import Any, override
+from software.evaluation.logs.log_interface import TimestampedEvalLog
+from software.evaluation.logs.world_state_log import WorldStateLog
+
+
+class EventType(StrEnum):
+    """Enum for the different types of events we want to track"""
+
+    PASS = auto()
+    SHOT_ON_GOAL = auto()
+    ENEMY_SHOT_ON_GOAL = auto()
+    SHOT_BLOCKED = auto()
+    FRIENDLY_POSSESSION_START = auto()
+    FRIENDLY_POSSESSION_END = auto()
+    ENEMY_POSSESSION_START = auto()
+    ENEMY_POSSESSION_END = auto()
+    GAME_START = auto()
+    GAME_END = auto()
+    GOAL_SCORED = auto()
+    YELLOW_CARD = auto()
+    RED_CARD = auto()
+
+
+class Team(StrEnum):
+    """The teams present in the game"""
+
+    BLUE = auto()
+    YELLOW = auto()
+
+
+@dataclass(kw_only=True)
+class EventLog(TimestampedEvalLog):
+    """Represents a single event being tracked, where and for whom the event is,
+    and the game state at the time of the event
+    """
+
+    event_type: EventType
+    from_team: Team
+    for_team: Team
+    world_state_log: WorldStateLog
+
+    num_cols = TimestampedEvalLog.get_num_cols() + 3 + WorldStateLog.get_num_cols()
+
+    @staticmethod
+    def from_world(
+        world_msg: World, event_type: EventType, from_team: Team, for_team: Team
+    ) -> EventLog:
+        """Creates an EventLog from a world protobuf message
+
+        :param world_msg: the world object containing the state of the game
+        :param event_type: the type of event being recorded
+        :param from_team: the team that the event is coming from
+        :param for_team: the team that the event is for
+        :return: a fully populated EventLog including world state
+        """
+        world_state_log = WorldStateLog.from_world(world_msg=world_msg)
+
+        return EventLog(
+            event_type=event_type,
+            from_team=from_team,
+            for_team=for_team,
+            world_state_log=world_state_log,
+        )
+
+    @override
+    @classmethod
+    def get_num_cols(cls) -> int:
+        return EventLog.num_cols
+
+    @override
+    def to_array(self) -> list[Any]:
+        return (
+            super().to_array()
+            + [
+                self.event_type.value,
+                self.from_team.value,
+                self.for_team.value,
+            ]
+            + self.world_state_log.to_array()
+        )
+
+    @staticmethod
+    @override
+    def from_csv_row(row_iter: Iterator[str]) -> EventLog | None:
+        """Parses a full CSV row into an EventLog."""
+        timestamp = float(next(row_iter))
+
+        event_type = EventType(next(row_iter))
+        from_team = Team(next(row_iter))
+        for_team = Team(next(row_iter))
+
+        world_state = WorldStateLog.from_csv_row(row_iter)
+
+        if not world_state:
+            return None
+
+        return EventLog(
+            timestamp=timestamp,
+            event_type=event_type,
+            from_team=from_team,
+            for_team=for_team,
+            world_state_log=world_state,
+        )