gameboy_worlds.emulation.tracker

   1from gameboy_worlds.emulation.parser import StateParser
   2from gameboy_worlds.utils import (
   3    nested_dict_to_str,
   4    verify_parameters,
   5    log_info,
   6    log_error,
   7    log_warn,
   8    log_dict,
   9    show_frames,
  10)
  11
  12
  13import numpy as np
  14from typing import Optional, Type, Dict, Any, Tuple, List, Set
  15
  16from abc import ABC, abstractmethod
  17
  18EPSILON = 0.001
  19""" Default epsilon for frame change detection. """
  20
  21
  22class MetricGroup(ABC):
  23    """
  24    Abstract Base class for organizing related metrics.
  25
  26    ### Documentation Guidlines:
  27    Every subchild should document the following in their class docstrings:
  28    - Reports (List of keys that are present in the return dict of `report`)
  29    - Final Reports (List of keys that are present in the return dict of `report_final`)
  30
  31    """
  32
  33    NAME = "base"
  34    """ Name of the MetricGroup. """
  35
  36    REQUIRED_PARSER = StateParser
  37    """ The StateParser which implements the minimum required functionality for this MetricGroup to work. """
  38
  39    def __init__(self, state_parser: StateParser, parameters: dict):
  40        verify_parameters(parameters)
  41        if not issubclass(type(state_parser), self.REQUIRED_PARSER):
  42            log_error(
  43                f"StateParser of type {type(state_parser)} is not compatible with MetricGroup requiring {self.REQUIRED_PARSER}."
  44            )
  45        self.state_parser = state_parser
  46        """ An instance of the StateParser to parse game state variables. """
  47        self._parameters = parameters
  48        self.start()
  49        self.final_metrics: Dict[str, Any] = None
  50        """ Dictionary to store final metrics after environment close. """
  51
  52    def start(self):
  53        """
  54        Called once when environment starts.
  55        All subclasses should call super() AFTER initializing their own variables.
  56        Only variables that will persist across episodes should be initialized here.
  57        """
  58        self.reset(first=True)
  59
  60    @abstractmethod
  61    def reset(self, first: bool = False):
  62        """Called when environment resets.
  63
  64        Args:
  65            first (bool): Whether this is the first reset of the environment. If True, might need to aggregate metrics into running final totals.
  66        """
  67        raise NotImplementedError
  68
  69    @abstractmethod
  70    def close(self):
  71        """
  72        Called when environment closes. Good for computing summary stats.
  73
  74        Step will not be called after this.
  75        """
  76        raise NotImplementedError
  77
  78    @abstractmethod
  79    def step(self, current_frame: np.ndarray, recent_frames: Optional[np.ndarray]):
  80        """
  81        Called each environment step to update metrics.
  82        Args:
  83            current_frame (np.ndarray): The current frame rendered by the emulator.
  84            recent_frames (Optional[np.ndarray]): The stack of frames that were rendered during the last action. Shape is [n_frames, height, width, channels]. Can be None if rendering is disabled.
  85        """
  86        raise NotImplementedError
  87
  88    @abstractmethod
  89    def report(self) -> Dict[str, Any]:
  90        """
  91        Return metrics as dictionary for instantaneous variable tracking.
  92
  93        :return: Dictionary of metrics
  94        :rtype: Dict[str, Any]
  95        """
  96        raise NotImplementedError
  97
  98    @abstractmethod
  99    def report_final(self) -> dict:
 100        """
 101        Return metrics as dictionary for logging. Called at end of environment (before close).
 102        Will never be called before `self.close`.
 103
 104        :return: Dictionary of metrics
 105        :rtype: Dict[str, Any]
 106        """
 107        raise NotImplementedError
 108
 109    def log_info(self, message: str):
 110        """
 111        Logs with MetricGroup's name. Primarily for debugging.
 112        """
 113        log_info(f"[Metric({self.NAME})]: {message}", self._parameters)
 114
 115    def log_warn(self, message: str):
 116        """
 117        Logs with MetricGroup's name. Primarily for debugging.
 118        """
 119        log_warn(f"[Metric({self.NAME})]: {message}", self._parameters)
 120
 121    def log_report(self):
 122        """
 123        Logs the current metrics report with MetricGroup's name. Primarily for debugging.
 124        """
 125        log_info(f"Metric({self.NAME}):\n")
 126        log_dict(self.report(), parameters=self._parameters)
 127
 128
 129class CoreMetrics(MetricGroup):
 130    """
 131    Tracks basic metrics that are guaranteed to be available in state tracker reports for any and all games:
 132
 133    Reports:
 134    - `steps`: Number of steps taken in the episode.
 135    - `frame_changed`: Whether the frame has changed since the last step.
 136    - `current_frame`: The current frame.
 137    - `passed_frames`: All frames that have passed since the last step.
 138
 139    Final Reports:
 140    - `total_episodes`: Total number of episodes completed.
 141    - `average_steps_per_episode`: Average number of steps taken per episode.
 142    - `max_steps`: Maximum number of steps taken in any episode.
 143    - `min_steps`: Minimum number of steps taken in any episode.
 144    - `std_steps`: Standard deviation of steps taken across episodes.
 145
 146    """
 147
 148    NAME = "core"
 149
 150    def start(self):
 151        self.steps_per_episode = []
 152        """ List of steps taken in each episode. """
 153        super().start()
 154
 155    def reset(self, first=False):
 156        if not first:
 157            self.steps_per_episode.append(self.steps)
 158        else:
 159            self.steps = 0
 160            """ Number of steps taken in the episode. """
 161            self.previous_frame = None
 162            """ Previous frame for detecting changes. """
 163            self.current_frame = None
 164            """ Current frame. """
 165            self.frame_changed = True
 166            """ Whether the frame has changed at all since last step. """
 167            self.passed_frames = None
 168            """ Stack of frames since the last step """
 169
 170    def close(self):
 171        if len(self.steps_per_episode) > 0:
 172            total_episodes = len(self.steps_per_episode)
 173            average_steps = np.mean(self.steps_per_episode)
 174            max_steps = np.max(self.steps_per_episode)
 175            min_steps = np.min(self.steps_per_episode)
 176            std_steps = np.std(self.steps_per_episode)
 177        else:
 178            total_episodes = 0
 179            average_steps = 0.0
 180            max_steps = 0
 181            min_steps = 0
 182            std_steps = 0.0
 183        self.final_metrics = {
 184            "total_episodes": int(total_episodes),
 185            "average_steps_per_episode": float(average_steps),
 186            "max_steps": int(max_steps),
 187            "min_steps": int(min_steps),
 188            "std_steps": float(std_steps),
 189        }
 190
 191    def step(self, current_frame, recent_frames):
 192        self.steps += 1
 193        self.current_frame = current_frame
 194        self.passed_frames = recent_frames
 195        if self.previous_frame is None:
 196            self.previous_frame = current_frame
 197            self.frame_changed = True
 198        else:
 199            frame_changed = False
 200            comparison_frame = self.previous_frame
 201            if recent_frames is None:
 202                recent_frames = np.array([current_frame])
 203            for frame in recent_frames:
 204                if np.abs(frame - comparison_frame).mean() > EPSILON:
 205                    frame_changed = True
 206                else:
 207                    frame_changed = False
 208                    comparison_frame = frame
 209                if frame_changed:
 210                    break
 211            self.frame_changed = frame_changed
 212        self.previous_frame = current_frame
 213
 214    def report(self):
 215        """
 216        Provides the following metrics:
 217        - `steps`: Number of steps taken in the episode.
 218        - `frame_changed`: Whether the frame has changed since the last step.
 219        - `current_frame`: The current frame.
 220        - `passed_frames`: All frames that have passed since the last step.
 221        """
 222        return {
 223            "steps": self.steps,
 224            "frame_changed": self.frame_changed,
 225            "current_frame": self.current_frame,
 226            "passed_frames": self.passed_frames,
 227        }
 228
 229    def report_final(self):
 230        """
 231        Provides the following metrics:
 232        - `total_episodes`: Total number of episodes completed.
 233        - `average_steps_per_episode`: Average number of steps taken per episode.
 234        - `max_steps`: Maximum number of steps taken in any episode.
 235        - `min_steps`: Minimum number of steps taken in any episode.
 236        - `std_steps`: Standard deviation of steps taken across episodes.
 237        """
 238        return self.final_metrics
 239
 240
 241class OCRegionMetric(MetricGroup, ABC):
 242    """
 243    Watch particular screen regions and capture subscreens for OCR when possible. Does not actually perform OCR itself, but makes it easy to capture the relevant regions.
 244    Children implementing this must define self.kinds in `start()` and then call on `super().start()`.
 245
 246    Reports:
 247    - `ocr_regions`: A dictionary mapping kinds to captured regions that had OCR-eligible text detected in them. The keys are kinds of OCR regions, and the values are the stacks of captured screen regions as numpy arrays of shape (num_captures, height, width, channels).
 248    - `step`: The current step number. Useful for differentiating when multiple OCR texts were found in the same episode. You can typically safely ignore this.
 249
 250    Final Reports:
 251    - `ocr_regions`: A list of tuples for all steps where OCR was detected. Is in form: `List[Tuple[int, Dict[str, np.ndarray]]]` where the int is the step number and the Dict maps kinds to a stack of the captured screen region.
 252
 253    """
 254
 255    NAME = "ocr"
 256
 257    def start(self):
 258        """
 259        Assumes the child has initialized a dict called self.kinds which tracks the various kinds of OCR that could be done.
 260                self.kinds should be in the form: {kind: region_name} where region_name is the name of the region to OCR for that kind.
 261                Will track ocr captured region results in form of list of dictionaries where these kinds are keys.
 262        """
 263        super().start()
 264        if self.NAME != "ocr":
 265            log_error(
 266                "OCRMetric subclasses must have NAME equal to 'ocr' for the environment get_info() aggregation step to work.",
 267                self._parameters,
 268            )
 269        if not hasattr(self, "kinds"):
 270            log_error("OCRMetrics must declare self.kinds dictionary", self._parameters)
 271        elif not isinstance(self.kinds, dict):
 272            log_error("self.kinds must be a dictionary", self._parameters)
 273        self.kinds: dict
 274        for item in self.kinds:
 275            if not isinstance(item, str):
 276                log_error("self.kinds keys must be strings", self._parameters)
 277            region_info = self.kinds[item]
 278            if not isinstance(region_info, str):
 279                log_error(
 280                    "self.kinds values must be region names (strings)", self._parameters
 281                )
 282            if region_info not in self.state_parser.named_screen_regions:
 283                log_error(
 284                    f"OCR region name {region_info} not found in state parser named regions. Available options: {self.state_parser.named_screen_regions}",
 285                    self._parameters,
 286                )
 287
 288    @staticmethod
 289    def can_read_kind(self, frame: np.ndarray, kind: str) -> bool:
 290        """
 291        Checks if the frame has text for the given kind.
 292
 293        Args:
 294            frame (np.ndarray): The frame to check.
 295            kind (str): The kind of text to check for.
 296        """
 297        raise NotImplementedError
 298
 299    def reset(self, first=False):
 300        """
 301        ocr_regions will track a list of the form List[Tuple[int, Dict[str, np.ndarray]]]
 302        which is a list of (step_number, {kind: ocr_region}) dictionaries.
 303        """
 304        self.ocr_regions = []
 305        self.steps = 0
 306        self.prev_has_ocr = False
 307
 308    def step(self, current_frame: np.ndarray, recent_frames: Optional[np.ndarray]):
 309        all_frames = None
 310        if recent_frames is not None:
 311            all_frames = recent_frames  # Current frame is included in recent frames
 312        else:
 313            all_frames = np.array([current_frame])
 314        ocr_dict = {}
 315        # Aggregate results for all frames and separate per kind.
 316        for kind in self.kinds.keys():
 317            captured_frames = []
 318            for frame in all_frames:
 319                if self.can_read_kind(frame, kind):
 320                    captured_frames.append(
 321                        self.state_parser.capture_named_region(
 322                            current_frame=frame, name=self.kinds[kind]
 323                        )
 324                    )
 325            true_captured_frames = []  # remove duplicates for efficiency
 326            # add the later frame in case of overlap detection. This prefers more text for cases where text comes as a stream.
 327            if len(captured_frames) > 0:
 328                for i in range(len(captured_frames) - 1):
 329                    curr_frame = captured_frames[i]
 330                    next_frame = captured_frames[i + 1]
 331                    if np.abs(curr_frame - next_frame).mean() > EPSILON:
 332                        true_captured_frames.append(curr_frame)
 333                true_captured_frames.append(
 334                    captured_frames[-1]
 335                )  # always add the last frame.
 336                ocr_dict[kind] = np.array(true_captured_frames)
 337        if len(ocr_dict) > 0:
 338            self.ocr_regions.append((self.steps, ocr_dict))
 339            self.prev_has_ocr = True
 340        else:
 341            self.prev_has_ocr = False
 342        self.steps += 1
 343
 344    def report(self):
 345        """
 346        Reports just the previous step's OCR regions if any were found.
 347        Returns:
 348            dict: A dictionary containing the captured regions for focused OCR.
 349        """
 350        if self.prev_has_ocr:
 351            return {
 352                "ocr_regions": self.ocr_regions[-1][1],
 353                "step": self.ocr_regions[-1][0],
 354            }
 355        else:
 356            return {}
 357
 358    def report_final(self):
 359        """
 360        Reports all the OCR regions extracted in the episode.
 361        """
 362        return {"ocr_regions": self.ocr_regions}
 363
 364    def close(self):
 365        pass
 366
 367
 368class SubGoal(ABC):
 369    """
 370    Abstract class representing a subgoal for tracking progress towards a test goal. These are intermediate states that must be achieved on the way to the final test goal.
 371    By convention, the task goal state itself is *not* considered a subgoal, but rather the final goal that the subgoals lead towards.
 372    """
 373
 374    NAME = "placeholder"
 375    """ Name of the subgoal. """
 376
 377    def __init__(self):
 378        if self.NAME == "placeholder":
 379            log_error(
 380                "Subclasses of SubGoal must set a unique NAME class variable.",
 381            )
 382        self.completed = False
 383
 384    def check_completed(self, frames: np.ndarray, parser: StateParser) -> bool:
 385        """
 386        Checks whether the subgoal has been completed based on the given frames and state parser.
 387
 388        Args:
 389            frames (np.ndarray): The stack of frames to check for subgoal completion.
 390            parser (StateParser): The state parser to use for checking subgoal completion.
 391        Returns:
 392            bool: True if the subgoal is completed, False otherwise.
 393        """
 394        for frame in frames:
 395            if self._check_completed(frame, parser):
 396                return True
 397        return False
 398
 399    @abstractmethod
 400    def _check_completed(self, frame: np.ndarray, parser: StateParser) -> bool:
 401        """
 402        Checks whether the subgoal has been completed based on a single frame and the state parser.
 403
 404        Args:
 405            frames (np.ndarray): A single frame to check for subgoal completion.
 406            parser (StateParser): The state parser to use for checking subgoal completion.
 407        Returns:
 408            bool: True if the subgoal is completed, False otherwise.
 409        """
 410        pass
 411
 412
 413class DummySubGoal(SubGoal):
 414    """
 415    A dummy subgoal that is never completed. Useful for testing.
 416    """
 417
 418    NAME = "dummy_subgoal"
 419
 420    def _check_completed(self, frame: np.ndarray, parser: StateParser) -> bool:
 421        return False
 422
 423
 424class RegionMatchSubGoal(SubGoal, ABC):
 425    """
 426    A subgoal that is completed if a specific region matches a target. Can be used to track subgoals that require specific dialogue boxes to appear, etc.
 427    """
 428
 429    NAME = "placeholder"
 430    _NAMED_REGION: str = None
 431    _TARGET_NAME: str = None
 432
 433    def __init__(self):
 434        super().__init__()
 435        if self._NAMED_REGION is None or self._TARGET_NAME is None:
 436            log_error(
 437                "Subclasses of RegionMatchSubGoal must set _NAMED_REGION and _TARGET_NAME class variables.",
 438            )
 439
 440    def _check_completed(self, frame: np.ndarray, parser: StateParser) -> bool:
 441        matches = parser.named_region_matches_multi_target(
 442            frame, self._NAMED_REGION, self._TARGET_NAME
 443        )
 444        return matches
 445
 446
 447class SingleRegionMatchSubGoal(SubGoal, ABC):
 448    """
 449    A subgoal that is completed if a specific single region matches its target.
 450    """
 451
 452    NAME = "placeholder"
 453    _NAMED_REGION: str = None
 454
 455    def __init__(self):
 456        super().__init__()
 457        if self._NAMED_REGION is None:
 458            log_error(
 459                "Subclasses of SingleRegionMatchSubGoal must set _NAMED_REGION class variable.",
 460            )
 461
 462    def _check_completed(self, frame, parser):
 463        matches = parser.named_region_matches_target(frame, self._NAMED_REGION)
 464        return matches
 465
 466
 467class AnyRegionMatchSubGoal(SubGoal, ABC):
 468    """
 469    A subgoal that is completed if any of a list of specific regions matches their targets.
 470    """
 471
 472    NAME = "placeholder"
 473    _NAMED_REGIONS: List[str] = None
 474    _TARGET_NAMES: List[str] = None
 475
 476    def __init__(self):
 477        super().__init__()
 478        if (
 479            self._NAMED_REGIONS is None
 480            or self._TARGET_NAMES is None
 481            or len(self._NAMED_REGIONS) != len(self._TARGET_NAMES)
 482            or len(self._NAMED_REGIONS) == 0
 483        ):
 484            log_error(
 485                "Subclasses of AnyRegionMatchSubGoal must set _NAMED_REGIONS and _TARGET_NAMES class variables, and they must be of the same length non zero.",
 486            )
 487
 488    def _check_completed(self, frame: np.ndarray, parser: StateParser) -> bool:
 489        for named_region, target_name in zip(self._NAMED_REGIONS, self._TARGET_NAMES):
 490            matches = parser.named_region_matches_multi_target(
 491                frame, named_region, target_name
 492            )
 493            if matches:
 494                return True
 495        return False
 496
 497
 498class AnySingleRegionMatchSubGoal(SubGoal, ABC):
 499    """
 500    A subgoal that is completed if any of a list of specific regions matches their targets, where each region only has one target.
 501    """
 502
 503    NAME = "placeholder"
 504    _NAMED_REGIONS: List[str] = None
 505
 506    def __init__(self):
 507        super().__init__()
 508        if self._NAMED_REGIONS is None or len(self._NAMED_REGIONS) == 0:
 509            log_error(
 510                "Subclasses of AnySingleRegionMatchSubGoal must set _NAMED_REGIONS class variable, and it must be non empty.",
 511            )
 512
 513    def _check_completed(self, frame: np.ndarray, parser: StateParser) -> bool:
 514        for named_region in self._NAMED_REGIONS:
 515            matches = parser.named_region_matches_target(frame, named_region)
 516            if matches:
 517                return True
 518        return False
 519
 520
 521class SubGoalMetric(MetricGroup, ABC):
 522    """
 523    Tracks subgoal based progress towards a specific test goal.
 524    Subgoals are always sequential, i.e. it is impossible to complete subgoal n+1 without completing subgoal n first.
 525
 526    Reports:
 527    - `all`: A list of the names of all subgoals being tracked, regardless of completion status.
 528    - `completed`: A list of the names of the subgoals that have been completed.
 529
 530    Final Reports:
 531    - `reached_subgoals`: List of subgoals that were reached at any point during any episode.
 532    """
 533
 534    NAME = "subgoals"
 535    SUBGOALS: List[SubGoal] = []
 536    """ List of SubGoal classes representing the subgoals to be tracked. These should be defined in child classes. """
 537
 538    def start(self):
 539        if self.NAME != "subgoals":
 540            log_error(
 541                f"SubGoalMetric NAME must be 'subgoals', got '{self.NAME}'.",
 542                self._parameters,
 543            )
 544        if len(self.SUBGOALS) == 0:
 545            log_error(
 546                "SubGoalMetric requires at least one subgoal to be defined in the SUBGOALS class variable.",
 547                self._parameters,
 548            )
 549        self._subgoals: List[SubGoal] = []
 550        """ List of SubGoal instances representing the subgoals being tracked. """
 551        for subgoal_class in self.SUBGOALS:
 552            subgoal_instance: SubGoal = subgoal_class()
 553            self._subgoals.append(subgoal_instance)
 554        self._reached_subgoals: Set[str] = set()
 555        """ Set of subgoals that were reached at any point during any episode. """
 556        super().start()
 557
 558    def close(self):
 559        pass
 560
 561    def reset(self, first=False):
 562        if not first:
 563            for subgoal in self._subgoals:
 564                if subgoal.completed:
 565                    self._reached_subgoals.add(subgoal.NAME)
 566        for subgoal in self._subgoals:
 567            subgoal.completed = False
 568
 569    def step(self, current_frame: np.ndarray, recent_frames: Optional[np.ndarray]):
 570        all_frames = None
 571        if recent_frames is not None:
 572            all_frames = recent_frames  # Current frame is included in recent frames
 573        else:
 574            all_frames = np.array([current_frame])
 575        for subgoal in self._subgoals:
 576            if not subgoal.completed:
 577                completed = subgoal.check_completed(all_frames, self.state_parser)
 578                subgoal.completed = completed
 579
 580    def report(self):
 581        """
 582        Reports the names of the subgoals being tracked and which ones have been completed.
 583        Returns:
 584            dict: A dictionary containing the list of all subgoals and the list of completed subgoals.
 585        """
 586        return {
 587            "all": [subgoal.NAME for subgoal in self._subgoals],
 588            "completed": [
 589                subgoal.NAME for subgoal in self._subgoals if subgoal.completed
 590            ],
 591        }
 592
 593    def report_final(self):
 594        """
 595        Reports the names of the subgoals that were reached at any point during any episode.
 596        Returns:
 597            dict: A dictionary containing the list of reached subgoals.
 598        """
 599        return {"reached_subgoals": list(self._reached_subgoals)}
 600
 601
 602class DummySubGoalMetric(SubGoalMetric):
 603    """
 604    A dummy SubGoalMetric that tracks a single DummySubGoal. Useful for testing.
 605    """
 606
 607    SUBGOALS = [DummySubGoal]
 608
 609
 610def make_subgoal_metric_class(subgoals: List[Type[SubGoal]]) -> Type[SubGoalMetric]:
 611    """
 612    Factory function to create a SubGoalMetric class with the given subgoals and name.
 613
 614    Args:
 615        subgoals (List[Type[SubGoal]]): The list of SubGoal classes to track.
 616        name (str): The name of the SubGoalMetric class.
 617
 618    Returns:
 619        Type[SubGoalMetric]: A new SubGoalMetric class with the specified subgoals and name.
 620    """
 621    if len(subgoals) == 0:
 622        log_error("Must provide at least one subgoal to create a SubGoalMetric class.")
 623
 624    class CustomSubGoalMetric(SubGoalMetric):
 625        SUBGOALS = subgoals
 626
 627    return CustomSubGoalMetric
 628
 629
 630class TerminationTruncationMetric(MetricGroup, ABC):
 631    """
 632    Tracks whether the environment was terminated or truncated.
 633
 634    Reports:
 635    - `terminated`: Whether the environment was terminated.
 636    - `truncated`: Whether the environment was truncated.
 637
 638    Final Reports:
 639    - `episode_end_reason`: List of reasons for episode endings: "terminated", "truncated", or None (None will occur only if there is a bug that leads to a premature reset).
 640    """
 641
 642    NAME = "termination_truncation"
 643
 644    def start(self):
 645        super().start()
 646        if self.NAME != "termination_truncation":
 647            log_error(
 648                f"TerminationTruncationMetric NAME must be 'termination_truncation', got '{self.NAME}'.",
 649                self._parameters,
 650            )
 651        self.episode_end_reason = []
 652        """ List of reasons for episode: termination or truncation or None (None will occur only if there is a bug that leads to a premature reset). """
 653        self.terminated = False
 654        """ Whether the environment was terminated. """
 655        self.truncated = False
 656        """ Whether the environment was truncated. """
 657
 658    def reset(self, first=False):
 659        if not first:
 660            if self.terminated:
 661                self.episode_end_reason.append("terminated")
 662            elif self.truncated:
 663                self.episode_end_reason.append("truncated")
 664            else:
 665                self.episode_end_reason.append(None)
 666        self.terminated = False
 667        self.truncated = False
 668
 669    def close(self):
 670        pass
 671
 672    @abstractmethod
 673    def determine_truncated(
 674        self, current_frame: np.ndarray, recent_frames: Optional[np.ndarray]
 675    ) -> bool:
 676        """
 677        Determines whether the environment was truncated.
 678
 679        :param current_frame: The current frame rendered by the emulator.
 680        :type current_frame: np.ndarray
 681        :param recent_frames: The stack of frames that were rendered during the last action. Shape is [n_frames, height, width, channels]. Can be None if rendering is disabled.
 682        :type recent_frames: Optional[np.ndarray]
 683        :return: True if the environment was truncated, False otherwise.
 684        :rtype: bool
 685        """
 686        pass
 687
 688    @abstractmethod
 689    def determine_terminated(
 690        self, current_frame: np.ndarray, recent_frames: Optional[np.ndarray]
 691    ) -> bool:
 692        """
 693        Determines whether the environment was terminated.
 694
 695        :param current_frame: The current frame rendered by the emulator.
 696        :type current_frame: np.ndarray
 697        :param recent_frames: The stack of frames that were rendered during the last action. Shape is [n_frames, height, width, channels]. Can be None if rendering is disabled.
 698        :type recent_frames: Optional[np.ndarray]
 699        :return: True if the environment was terminated, False otherwise.
 700        :rtype: bool
 701        """
 702        pass
 703
 704    def step(self, current_frame: np.ndarray, recent_frames: Optional[np.ndarray]):
 705        """
 706        Determines whether the environment was terminated or truncated.
 707        """
 708        if self.terminated or self.truncated:
 709            return  # This should ideally not happen, because the environment should reset after termination or truncation.
 710        self.truncated = self.determine_truncated(current_frame, recent_frames)
 711        self.terminated = self.determine_terminated(current_frame, recent_frames)
 712
 713    def report(self):
 714        """
 715        Reports whether the environment was terminated or truncated.
 716        Returns:
 717            dict: A dictionary containing the termination and truncation status.
 718        """
 719        return {
 720            "terminated": self.terminated,
 721            "truncated": self.truncated,
 722        }
 723
 724    def report_final(self):
 725        """
 726        Reports the reasons for episode endings.
 727        Returns:
 728            dict: A dictionary containing the list of episode end reasons.
 729        """
 730        return {"episode_end_reason": self.episode_end_reason}
 731
 732
 733class TerminationMetric(TerminationTruncationMetric, ABC):
 734    def determine_truncated(
 735        self, current_frame: np.ndarray, recent_frames: Optional[np.ndarray]
 736    ) -> bool:
 737        return False
 738
 739
 740class StateTracker:
 741    """
 742    Tracks and provides API access to the game state / metrics over time and across episodes.
 743    The most hassle-free way to read from the StateTracker is to use the `report()` and `report_final()` methods to get nested dictionaries of all metrics tracked.
 744
 745    **Example Usage:**
 746
 747    ```python
 748    import numpy as np
 749    from gameboy_worlds import get_pokemon_emulator
 750    emulator = get_pokemon_emulator(variant="pokemon_red")
 751
 752    # We can access the StateTracker via the emulator
 753    state_tracker = emulator.state_tracker
 754
 755    # Run a random action on the emulator
 756    emulator.reset()
 757    allowed_actions = list(LowLevelActions)
 758    action = np.random.choice(allowed_actions)
 759    _, _ = emulator.step(action) # also updates the StateTracker internally
 760    # We can access the current episode metrics via the StateTracker
 761    episode_metrics = state_tracker.report() # access all of them as a nested dict
 762    specific_metric = state_tracker.get_episode_metric(("core", "steps")) # access specific metrics
 763
 764    # If we reset the emulator, the StateTracker will reset its inter-episode metrics as well
 765    emulator.reset()
 766    action = np.random.choice(allowed_actions)
 767    _, _ = emulator.step(action)
 768    emulator.close() # StateTracker will finalize its metrics internally
 769    final_metrics = state_tracker.report_final() # access all of them as a nested dict
 770    specific_final_metric = state_tracker.get_final_metric(("core", "average_steps_per_episode")) # access specific final metrics
 771    ```
 772    """
 773
 774    TERMINATION_TRUNCATION_METRIC: Type[TerminationTruncationMetric] = None
 775    """ The TerminationTruncationMetric class to use for tracking termination and truncation. If None, no such metric will be tracked. """
 776
 777    SUBGOAL_METRIC: Type[SubGoalMetric] = None
 778    """ The SubGoalMetric class to use for tracking subgoal progress. If None, no such metric will be tracked. """
 779
 780    def __init__(
 781        self,
 782        state_parser: StateParser,
 783        parameters: dict,
 784    ):
 785        """
 786        Initializes the StateTracker.
 787        Args:
 788            state_parser (StateParser): An instance of the StateParser to parse game state variables.
 789            parameters (dict): A dictionary of parameters for configuration.
 790        """
 791        verify_parameters(parameters)
 792        self.state_parser = state_parser
 793        """ An instance of the StateParser to parse game state variables. """
 794        self._parameters = parameters
 795        self.start()
 796        self.validate()
 797        if self.metric_classes[0] != CoreMetrics:
 798            log_error(
 799                "First metric class must be CoreMetrics. Make sure to call `super().start()` first in child class overrides of `start()`.",
 800                parameters,
 801            )
 802        self.metrics = {}
 803        """ Dictionary to store MetricGroup instances. """
 804        for metric_group_class in self.metric_classes:
 805            metric_group_instance: MetricGroup = metric_group_class(
 806                state_parser, parameters
 807            )
 808            self.metrics[metric_group_instance.NAME] = metric_group_instance
 809        self.episode_metrics: Dict[str, Dict[str, Any]] = {}
 810        """ Dictionary to store metrics running during episode. """
 811        self.final_metrics: Dict[str, Dict[str, Any]] = {}
 812
 813    def start(self):
 814        """
 815        Sets up the metrics for the tracker by creating the list `self.metric_classes`
 816
 817        Child classes must FIRST call super().start() and THEN set up their own metric classes.
 818        """
 819        self.metric_classes: List[Type[MetricGroup]] = [CoreMetrics]
 820        if self.TERMINATION_TRUNCATION_METRIC is not None:
 821            if not issubclass(
 822                self.TERMINATION_TRUNCATION_METRIC, TerminationTruncationMetric
 823            ):
 824                log_error(
 825                    "TERMINATION_TRUNCATION_METRIC must be a subclass of TerminationTruncationMetric.",
 826                    self._parameters,
 827                )
 828            self.metric_classes.append(self.TERMINATION_TRUNCATION_METRIC)
 829        if self.SUBGOAL_METRIC is not None:
 830            if not issubclass(self.SUBGOAL_METRIC, SubGoalMetric):
 831                log_error(
 832                    "SUBGOAL_METRIC must be a subclass of SubGoalMetric.",
 833                    self._parameters,
 834                )
 835            self.metric_classes.append(self.SUBGOAL_METRIC)
 836
 837    def validate(self):
 838        """
 839        Is meant to be called once after initialization to ensure that the tracker is valid.
 840        """
 841        pass
 842
 843    def reset(self):
 844        """
 845        Is called once per environment reset to reset any tracked metrics.
 846        """
 847        for metric_group in self.metrics.values():
 848            metric_group.reset()
 849        self.step()
 850
 851    def step(self, recent_frames: Optional[np.ndarray] = None):
 852        """
 853        Is called once per environment step to update any tracked metrics.
 854
 855        Args:
 856            recent_frames (Optional[np.ndarray]): The stack of frames that were rendered during the last action. Shape is [n_frames, height, width, channels]. Can be None if rendering is disabled.
 857            epsilon (float, optional): The threshold for considering a frame change.
 858        """
 859        current_frame = None
 860        if recent_frames is None:
 861            current_frame = self.state_parser.get_current_frame()
 862        else:
 863            current_frame = recent_frames[-1]
 864        self.episode_metrics = {}
 865        for metric_group in self.metrics.values():
 866            metric_group.step(current_frame, recent_frames)
 867            self.episode_metrics[metric_group.NAME] = metric_group.report()
 868
 869    def close(self):
 870        """
 871        Is called once when the environment is closed to finalize any tracked metrics.
 872        """
 873        for metric_group in self.metrics.values():
 874            metric_group.close()
 875        self.final_metrics = {
 876            name: mg.report_final() for name, mg in self.metrics.items()
 877        }
 878
 879    def report(self) -> Dict[str, Dict[str, Any]]:
 880        """
 881        Returns the current episode metrics.
 882
 883        :return: A nested dictionary containing the current episode metrics.
 884        :rtype: Dict[str, Dict[str, Any]]
 885        """
 886        return self.episode_metrics
 887
 888    def report_final(self) -> Dict[str, Dict[str, Any]]:
 889        """
 890        Returns the final metrics after environment close.
 891
 892        Returns:
 893            Dict[str, Dict[str, Any]]: A nested dictionary containing the final metrics.
 894        """
 895        return self.final_metrics
 896
 897    def _get_specific_metric(self, metrics_dict, key: Tuple[str, str]):
 898        if metrics_dict is None:
 899            log_error("No metrics available. Have you called step() or close()?")
 900        metric_group_name, metric_name = key
 901        if metric_group_name not in metrics_dict:
 902            log_error(
 903                f"Metric group {metric_group_name} not found in metrics. Available groups: {list(metrics_dict.keys())}"
 904            )
 905        if metric_name not in metrics_dict[metric_group_name]:
 906            log_error(
 907                f"Metric {metric_name} not found in metric group {metric_group_name}. Available metrics: {list(metrics_dict[metric_group_name].keys())}"
 908            )
 909        return metrics_dict[metric_group_name][metric_name]
 910
 911    def get_episode_metric(self, key: Tuple[str, str]):
 912        """
 913        Returns the metrics for a specific episode and metric group.
 914
 915        Does not give final metrics at any point.
 916
 917        :param key: A tuple of the form (metric_group_name, metric_name).
 918        :type key: Tuple[str, str]
 919        :return: The requested metric value
 920        :rtype: Any
 921        """
 922        return self._get_specific_metric(self.episode_metrics, key)
 923
 924    def get_final_metric(self, key: Tuple[str, str]):
 925        """
 926        Returns the final metrics for a specific metric group.
 927
 928        :param key: A tuple of the form (metric_group_name, metric_name).
 929        :type key: Tuple[str, str]
 930        :return: The requested final metric value
 931        :rtype: Any
 932        """
 933        return self._get_specific_metric(self.final_metrics, key)
 934
 935    def __repr__(self) -> str:
 936        metric_names = [mg.NAME for mg in self.metrics.values()]
 937        return f"<StateTracker, metrics=({', '.join(metric_names)})>"
 938
 939
 940class TestTrackerMixin:
 941    """
 942    Mixin class for testing trackers.
 943    Ensures that exactly one of the tracked metrics is a TerminationTruncationMetric.
 944    """
 945
 946    def validate(self):
 947        if not hasattr(self, "_parameters"):
 948            log_error("Parameters have not been set yet.")
 949        if self.TERMINATION_TRUNCATION_METRIC is None:
 950            log_error(
 951                "TestTrackerMixin requires a TerminationTruncationMetric to be set as TERMINATION_TRUNCATION_METRIC.",
 952                self._parameters,
 953            )
 954        if self.SUBGOAL_METRIC is None:
 955            log_error(
 956                "TestTrackerMixin requires a SubGoalMetric to be set as SUBGOAL_METRIC.",
 957                self._parameters,
 958            )
 959
 960
 961class RegionMatchTruncationMetric(TerminationTruncationMetric, ABC):
 962    """
 963    Truncates the episode if a specific region matches a target.
 964    Can be used to truncate episodes when specific dialogue boxes appear, etc.
 965    """
 966
 967    _TRUNCATION_NAMED_REGION = None
 968    _TRUNCATION_TARGET_NAME = None
 969
 970    def determine_truncated(self, current_frame, recent_frames):
 971        if (
 972            self._TRUNCATION_NAMED_REGION is None
 973            or self._TRUNCATION_TARGET_NAME is None
 974        ):
 975            log_error(
 976                "Must set _TRUNCATION_NAMED_REGION and _TRUNCATION_TARGET_NAME.",
 977                self._parameters,
 978            )
 979        all_frames = [current_frame]
 980        if recent_frames is not None:
 981            all_frames = recent_frames
 982        for frame in all_frames:
 983            matches = self.state_parser.named_region_matches_multi_target(
 984                frame,
 985                self._TRUNCATION_NAMED_REGION,
 986                self._TRUNCATION_TARGET_NAME,
 987            )
 988            if matches:
 989                return True
 990        return False
 991
 992
 993class RegionMatchTerminationMetric(TerminationTruncationMetric, ABC):
 994    """
 995    Terminates the episode if a specific region matches a target.
 996    Can be used to terminate episodes when specific dialogue boxes appear, etc.
 997    """
 998
 999    _TERMINATION_NAMED_REGION = None
1000    _TERMINATION_TARGET_NAME = None
1001
1002    def determine_terminated(self, current_frame, recent_frames):
1003        if (
1004            self._TERMINATION_NAMED_REGION is None
1005            or self._TERMINATION_TARGET_NAME is None
1006        ):
1007            log_error(
1008                "Must set _TERMINATION_NAMED_REGION and _TERMINATION_TARGET_NAME.",
1009                self._parameters,
1010            )
1011        all_frames = [current_frame]
1012        if recent_frames is not None:
1013            all_frames = recent_frames
1014        for frame in all_frames:
1015            matches = self.state_parser.named_region_matches_multi_target(
1016                frame,
1017                self._TERMINATION_NAMED_REGION,
1018                self._TERMINATION_TARGET_NAME,
1019            )
1020            if matches:
1021                return True
1022        return False
1023
1024
1025class RegionMatchTerminationOnlyMetric(TerminationMetric, ABC):
1026    """
1027    RegionMatchTerminationMetric with no truncation.
1028    No truncation.
1029    """
1030
1031    _TERMINATION_NAMED_REGION = None
1032    _TERMINATION_TARGET_NAME = None
1033
1034    def determine_terminated(self, current_frame, recent_frames):
1035        if (
1036            self._TERMINATION_NAMED_REGION is None
1037            or self._TERMINATION_TARGET_NAME is None
1038        ):
1039            log_error(
1040                "Must set _TERMINATION_NAMED_REGION and _TERMINATION_TARGET_NAME.",
1041                self._parameters,
1042            )
1043        all_frames = [current_frame]
1044        if recent_frames is not None:
1045            all_frames = recent_frames
1046        for frame in all_frames:
1047            matches = self.state_parser.named_region_matches_multi_target(
1048                frame,
1049                self._TERMINATION_NAMED_REGION,
1050                self._TERMINATION_TARGET_NAME,
1051            )
1052            if matches:
1053                return True
1054        return False
1055
1056class AnyRegionMatchTerminationMetric(TerminationMetric, ABC):
1057    """
1058    Terminates the episode if any of a list of specific regions matches their targets.
1059    No truncation.
1060    """
1061
1062    _NAMED_REGIONS: List[str] = None
1063    _TARGET_NAMES: List[str] = None
1064
1065    def __init__(self):
1066        super().__init__()
1067        if (
1068            self._NAMED_REGIONS is None
1069            or self._TARGET_NAMES is None
1070            or len(self._NAMED_REGIONS) != len(self._TARGET_NAMES)
1071            or len(self._NAMED_REGIONS) == 0
1072        ):
1073            log_error(
1074                "Subclasses of AnyRegionMatchTerminationMetric must set _NAMED_REGIONS and _TARGET_NAMES class variables, and they must be of the same length non zero.",
1075            )
1076
1077    def determine_terminated(self, current_frame, recent_frames):
1078        all_frames = [current_frame]
1079        if recent_frames is not None:
1080            all_frames = recent_frames
1081        for frame in all_frames:
1082            for named_region, target_name in zip(self._NAMED_REGIONS, self._TARGET_NAMES):
1083                if self.state_parser.named_region_matches_multi_target(
1084                    frame, named_region, target_name
1085                ):
1086                    return True
1087        return False
1088
1089
1090class RegionChangedTerminationMetric(TerminationTruncationMetric, ABC):
1091    """
1092    Terminates the episode if a specific named region changes significantly
1093    from its appearance at the start of the episode (on reset).
1094
1095    Useful for detecting pickups, stat changes, or any event that alters a
1096    HUD region without having a fixed reference capture.
1097
1098    Subclass and set:
1099        _CHANGED_NAMED_REGION: name of the NamedScreenRegion to monitor
1100        _CHANGE_MAE_THRESHOLD: MAE threshold above which the region is considered changed (default 10)
1101    """
1102
1103    _CHANGED_NAMED_REGION = None
1104    _CHANGE_MAE_THRESHOLD = 10
1105
1106    def reset(self, first=False):
1107        super().reset(first=first)
1108        self._region_baseline = None
1109
1110    def determine_truncated(self, current_frame, recent_frames):
1111        return False
1112
1113    def determine_terminated(self, current_frame, recent_frames):
1114        if self._CHANGED_NAMED_REGION is None:
1115            log_error("Must set _CHANGED_NAMED_REGION.", self._parameters)
1116        cropped = self.state_parser.capture_named_region(
1117            current_frame, self._CHANGED_NAMED_REGION
1118        )
1119        if self._region_baseline is None:
1120            self._region_baseline = cropped.copy()
1121            return False
1122        mae = np.abs(cropped.astype(float) - self._region_baseline.astype(float)).mean()
1123        return mae > self._CHANGE_MAE_THRESHOLD