gameboy_worlds.emulation.pokemon.trackers

  1from gameboy_worlds.emulation.pokemon.base_metrics import (
  2    CorePokemonMetrics,
  3    PokemonOCRMetric,
  4    PokemonRedLocation,
  5    PokemonRedStarter,
  6)
  7from gameboy_worlds.emulation.pokemon.test_metrics import (
  8    PokemonCenterTerminateMetric,
  9    OutsideViridianCenterSubgoal,
 10    MtMoonTerminateMetric,
 11    SpeakToBillCompleteTerminateMetric,
 12    PickupPokeballTerminateMetric,
 13    ReadTrainersTipsSignTerminateMetric,
 14    SpeakToCinnabarGymAideCompleteTerminateMetric,
 15    SpeakToCinnabarMonkTerminateMetric,
 16    DefeatedBrockTerminateMetric,
 17    DefeatedLassTerminateMetric,
 18    CaughtPidgeyTerminateMetric,
 19    CaughtPikachuTerminateMetric,
 20    BoughtPotionAtPewterPokemartTerminateMetric,
 21    UsedPotionOnCharmanderTerminateMetric,
 22    OpenMapTerminateMetric,
 23    PokemonPrismFirstBadgeTerminateMetric,
 24)
 25
 26from gameboy_worlds.emulation.pokemon.base_metrics import (
 27    PokemonTestMetric,
 28)
 29from gameboy_worlds.utils import log_info
 30from gameboy_worlds.emulation.tracker import (
 31    StateTracker,
 32    TestTrackerMixin,
 33    DummySubGoalMetric,
 34    make_subgoal_metric_class,
 35)
 36from gameboy_worlds.emulation.pokemon.parsers import (
 37    AgentState,
 38)
 39from typing import Optional
 40
 41
 42class CorePokemonTracker(StateTracker):
 43    """
 44    StateTracker for core Pokémon metrics.
 45    """
 46
 47    _ADD_GRID_OVERLAY = False
 48    """ Whether to add the grid overlay drawn by the state parser when the agent is in FREE ROAM. This is useful for VLM based agents may need a coordinate grid overlayed onto the frame, but may cause issues for agents that do not understand that it is not a part of the game. """
 49
 50    def start(self):
 51        super().start()
 52        self.metric_classes.extend([CorePokemonMetrics, PokemonTestMetric])
 53
 54    def step(self, *args, **kwargs):
 55        """
 56        Calls on super().step(), but then modifies the current frame to overlay the grid if the agent is in FREE ROAM.
 57        """
 58        super().step(*args, **kwargs)
 59        if self._ADD_GRID_OVERLAY:
 60            state = self.episode_metrics["pokemon_core"]["agent_state"]
 61            # if agent_state is in FREE ROAM, draw the grid, otherwise do not
 62            if state == AgentState.FREE_ROAM:
 63                screen = self.episode_metrics["core"]["current_frame"]
 64                screen = self.state_parser.draw_grid_overlay(current_frame=screen)
 65                self.episode_metrics["core"]["current_frame"] = screen
 66                previous_screens = self.episode_metrics["core"]["passed_frames"]
 67                if previous_screens is not None:
 68                    self.episode_metrics["core"]["passed_frames"][-1, :] = screen
 69
 70
 71class PokemonOCRTracker(CorePokemonTracker):
 72    def start(self):
 73        super().start()
 74        self.metric_classes.extend([PokemonOCRMetric])
 75
 76
 77class PokemonRedStarterTracker(PokemonOCRTracker):
 78    """
 79    Example StateTracker that tracks the starter Pokémon chosen in Pokémon Red.
 80    """
 81
 82    def start(self):
 83        super().start()
 84        self.metric_classes.extend([PokemonRedStarter, PokemonRedLocation])
 85
 86
 87class PokemonTestTracker(TestTrackerMixin, PokemonOCRTracker):
 88    """
 89    Inherit this class and set TERMINATION_TRUNCATION_METRIC to create a TestTracker for Pokémon games.
 90    """
 91
 92    TERMINATION_TRUNCATION_METRIC = PokemonCenterTerminateMetric
 93    SUBGOAL_METRIC = DummySubGoalMetric
 94
 95
 96class PokemonRedCenterTestTracker(PokemonTestTracker):
 97    """
 98    A TestTracker for Pokémon Red that ends an episode when the agent reaches the entrance to the Viridian City Pokémon Center.
 99    """
100
101    TERMINATION_TRUNCATION_METRIC = PokemonCenterTerminateMetric
102    SUBGOAL_METRIC = make_subgoal_metric_class([OutsideViridianCenterSubgoal])
103
104
105class PokemonRedMtMoonTestTracker(PokemonTestTracker):
106    """
107    A TestTracker for Pokémon Red that ends an episode when the agent reaches the entrance to Mt. Moon.
108    """
109
110    TERMINATION_TRUNCATION_METRIC = MtMoonTerminateMetric
111    SUBGOAL_METRIC = DummySubGoalMetric
112
113
114class PokemonRedSpeakToBillTestTracker(PokemonTestTracker):
115    """
116    A TestTracker for Pokémon Red that ends an episode when the agent speaks to Bill.
117    """
118
119    TERMINATION_TRUNCATION_METRIC = SpeakToBillCompleteTerminateMetric
120    SUBGOAL_METRIC = DummySubGoalMetric
121
122
123class PokemonRedPickupPokeballTestTracker(PokemonTestTracker):
124    """
125    A TestTracker for Pokémon Red that ends an episode when the agent picks up the Pokéball in Professor Oak's Lab.
126    """
127
128    TERMINATION_TRUNCATION_METRIC = PickupPokeballTerminateMetric
129    SUBGOAL_METRIC = DummySubGoalMetric
130
131
132class PokemonRedReadTrainersTipsSignTestTracker(PokemonTestTracker):
133    """
134    A TestTracker for Pokémon Red that ends an episode when the agent reads the Trainer's Tips sign.
135    """
136
137    TERMINATION_TRUNCATION_METRIC = ReadTrainersTipsSignTerminateMetric
138    SUBGOAL_METRIC = DummySubGoalMetric
139
140
141class PokemonRedSpeakToCinnabarGymAideCompleteTestTracker(PokemonTestTracker):
142    """
143    A TestTracker for Pokémon Red that ends an episode when the agent speaks to the Cinnabar Gym aide.
144    """
145
146    TERMINATION_TRUNCATION_METRIC = SpeakToCinnabarGymAideCompleteTerminateMetric
147    SUBGOAL_METRIC = DummySubGoalMetric
148
149
150class PokemonRedSpeakToCinnabarMonkTestTracker(PokemonTestTracker):
151    """
152    A TestTracker for Pokémon Red that ends an episode when the agent speaks to the Cinnabar Monk.
153    """
154
155    TERMINATION_TRUNCATION_METRIC = SpeakToCinnabarMonkTerminateMetric
156    SUBGOAL_METRIC = DummySubGoalMetric
157
158
159class PokemonRedDefeatedBrockTestTracker(PokemonTestTracker):
160    """
161    A TestTracker for Pokémon Red that ends an episode when the agent defeats Brock.
162    """
163
164    TERMINATION_TRUNCATION_METRIC = DefeatedBrockTerminateMetric
165    SUBGOAL_METRIC = DummySubGoalMetric
166
167
168class PokemonRedDefeatedLassTestTracker(PokemonTestTracker):
169    """
170    A TestTracker for Pokémon Red that ends an episode when the agent defeats the Lass trainer.
171    """
172
173    TERMINATION_TRUNCATION_METRIC = DefeatedLassTerminateMetric
174    SUBGOAL_METRIC = DummySubGoalMetric
175
176
177class PokemonRedCaughtPidgeyTestTracker(PokemonTestTracker):
178    """
179    A TestTracker for Pokémon Red that ends an episode when the agent catches a Pidgey.
180    """
181
182    TERMINATION_TRUNCATION_METRIC = CaughtPidgeyTerminateMetric
183    SUBGOAL_METRIC = DummySubGoalMetric
184
185
186class PokemonRedCaughtPikachuTestTracker(PokemonTestTracker):
187    """
188    A TestTracker for Pokémon Red that ends an episode when the agent catches a Pikachu.
189    """
190
191    TERMINATION_TRUNCATION_METRIC = CaughtPikachuTerminateMetric
192    SUBGOAL_METRIC = DummySubGoalMetric
193
194
195class PokemonRedBoughtPotionAtPewterPokemartTestTracker(PokemonTestTracker):
196    """
197    A TestTracker for Pokémon Red that ends an episode when the agent buys a Potion at the Pewter City Poké Mart.
198    """
199
200    TERMINATION_TRUNCATION_METRIC = BoughtPotionAtPewterPokemartTerminateMetric
201    SUBGOAL_METRIC = DummySubGoalMetric
202
203
204class PokemonRedUsedPotionOnCharmanderTestTracker(PokemonTestTracker):
205    """
206    A TestTracker for Pokémon Red that ends an episode when the agent uses a Potion on Charmander.
207    """
208
209    TERMINATION_TRUNCATION_METRIC = UsedPotionOnCharmanderTerminateMetric
210    SUBGOAL_METRIC = DummySubGoalMetric
211
212
213class PokemonRedOpenMapTestTracker(PokemonTestTracker):
214    """
215    A TestTracker for Pokémon Red that ends an episode when the agent opens the map.
216    """
217
218    TERMINATION_TRUNCATION_METRIC = OpenMapTerminateMetric
219    SUBGOAL_METRIC = DummySubGoalMetric
220
221
222class PokemonPrismFirstBadgeTestTracker(PokemonTestTracker):
223    """
224    TestTracker for Pokemon Prism: terminates when the player earns the first Naljo badge
225    (Magma Badge from Gym Leader Tansy in Brimstone City).
226    Truncates if the agent exits a battle without winning.
227    """
228
229    TERMINATION_TRUNCATION_METRIC = PokemonPrismFirstBadgeTerminateMetric
230    SUBGOAL_METRIC = DummySubGoalMetric