gameboy_worlds.emulation.registry

  1"""
  2Keeps a record of:
  3- Available games
  4- Expected save file names for each game
  5- Strongest available `State Parser`s for each game
  6- Available `State Tracker`s for each game, with string identifiers and a default tracker for each game
  7- Available `Emulator`s for each game, with string identifiers and a default emulator for each game
  8
  9Provides methods to access these.
 10"""
 11
 12from gameboy_worlds.utils import (
 13    log_error,
 14    load_parameters,
 15    log_warn,
 16    get_benchmark_tasks,
 17)
 18import os
 19from typing import Optional, Union, Type, Dict
 20from gameboy_worlds.emulation.parser import StateParser, DummyParser
 21from gameboy_worlds.emulation.tracker import StateTracker
 22from gameboy_worlds.emulation.emulator import Emulator
 23
 24from gameboy_worlds.emulation.pokemon import registry as pokemon_registry
 25from gameboy_worlds.emulation.legend_of_zelda import (
 26    registry as legend_of_zelda_registry,
 27)
 28from gameboy_worlds.emulation.hamtaro import registry as hamtaro_registry
 29from gameboy_worlds.emulation.sword_of_hope import registry as sword_of_hope_registry
 30from gameboy_worlds.emulation.deja_vu import registry as deja_vu_registry
 31from gameboy_worlds.emulation.harvest_moon import registry as harvest_moon_registry
 32from gameboy_worlds.emulation.runes_of_virtue import registry as runes_of_virtue_registry
 33from gameboy_worlds.emulation.harry_potter import registry as harry_potter_registry
 34from gameboy_worlds.emulation.bomberman import registry as bomberman_registry
 35from gameboy_worlds.emulation.survival_kids import registry as survival_kids_registry
 36
 37_game_registries = [
 38    hamtaro_registry,
 39    pokemon_registry,
 40    legend_of_zelda_registry,
 41    sword_of_hope_registry,
 42    deja_vu_registry,
 43    harvest_moon_registry,
 44    survival_kids_registry,
 45    runes_of_virtue_registry,
 46    harry_potter_registry,
 47    bomberman_registry,
 48]
 49
 50_project_parameters = load_parameters()
 51
 52
 53def _merge_into_dict(destination: dict, incoming: dict):
 54    """
 55    Merges incoming into destination, throwing an error if there are any overlapping keys.
 56
 57    Args:
 58        destination (dict): The destination dictionary.
 59        incoming (dict): The incoming dictionary.
 60
 61    Returns:
 62        dict: The merged dictionary.
 63    """
 64    for key in incoming:
 65        if key in destination:
 66            log_error(
 67                f"Duplicate key '{key}' found when merging dictionaries. This likely means there is a duplicate entry for '{key}' in the registry. Please check the registry for duplicates and remove them.\nDictionaries being merged:\nDestination: {destination}\nIncoming: {incoming}",
 68                _project_parameters,
 69            )
 70    destination.update(incoming)
 71
 72
 73GAME_TO_GB_NAME: Dict[str, str] = {}
 74""" Expected save name for each game. Save the file to <storage_dir_from_config_file>/<game_name>_rom_data/<gb_name>"""
 75
 76_STRONGEST_PARSERS: Dict[str, Type[StateParser]] = {}
 77""" Mapping of game names to their corresponding strongest StateParser classes. 
 78Unless you have a very good reason, you should always use the STRONGEST possible parser for a given game. 
 79The parser itself does not affect performance, as for it to perform a read / screen comparison operation , it must be called upon by the state tracker.
 80This means there is never a reason to use a weaker parser. 
 81"""
 82
 83AVAILABLE_STATE_TRACKERS: Dict[str, Dict[str, Type[StateTracker]]] = {}
 84""" Mapping of game names to their available StateTracker classes with string identifiers. """
 85
 86AVAILABLE_EMULATORS: Dict[str, Dict[str, Type[Emulator]]] = {}
 87""" Mapping of game names to their available Emulator classes with string identifiers. """
 88
 89for module in _game_registries:
 90    if hasattr(module, "GAME_TO_GB_NAME"):
 91        _merge_into_dict(GAME_TO_GB_NAME, module.GAME_TO_GB_NAME)
 92    else:
 93        log_error(
 94            f"Module '{module.__name__}' is missing the GAME_TO_GB_NAME mapping. ",
 95            _project_parameters,
 96        )
 97    if hasattr(module, "STRONGEST_PARSERS"):
 98        _merge_into_dict(_STRONGEST_PARSERS, module.STRONGEST_PARSERS)
 99    else:
100        log_error(
101            f"Module '{module.__name__}' is missing the STRONGEST_PARSERS mapping. This mapping is required for the registry to function properly.",
102            _project_parameters,
103        )
104    if hasattr(module, "AVAILABLE_STATE_TRACKERS"):
105        _merge_into_dict(AVAILABLE_STATE_TRACKERS, module.AVAILABLE_STATE_TRACKERS)
106    else:
107        log_error(
108            f"Module '{module.__name__}' is missing the AVAILABLE_STATE_TRACKERS mapping. This mapping is required for the registry to function properly.",
109            _project_parameters,
110        )
111    if hasattr(module, "AVAILABLE_EMULATORS"):
112        _merge_into_dict(AVAILABLE_EMULATORS, module.AVAILABLE_EMULATORS)
113    else:
114        log_error(
115            f"Module '{module.__name__}' is missing the AVAILABLE_EMULATORS mapping. This mapping is required for the registry to function properly.",
116            _project_parameters,
117        )
118
119
120for game in AVAILABLE_STATE_TRACKERS:
121    if "default" not in AVAILABLE_STATE_TRACKERS[game]:
122        log_error(
123            f"Game '{game}' is missing a default StateTracker mapping in the registry.",
124            _project_parameters,
125        )
126
127for game in AVAILABLE_EMULATORS:
128    if "default" not in AVAILABLE_EMULATORS[game]:
129        log_error(
130            f"Game '{game}' is missing a default Emulator mapping in the registry.",
131            _project_parameters,
132        )
133
134AVAILABLE_GAMES = list(GAME_TO_GB_NAME.keys())
135""" List of available games. """
136
137for game in AVAILABLE_GAMES:
138    if game not in _STRONGEST_PARSERS:
139        if _project_parameters["debug_mode"]:
140            log_warn(
141                f"Warning: Game '{game}' is missing a strongest StateParser mapping in the registry.",
142                _project_parameters,
143            )
144        else:
145            log_error(
146                f"Game '{game}' is missing a strongest StateParser mapping in the registry.",
147                _project_parameters,
148            )
149    if game not in AVAILABLE_STATE_TRACKERS:
150        if _project_parameters["debug_mode"]:
151            log_warn(
152                f"Warning: Game '{game}' is missing a StateTracker mapping in the registry.",
153                _project_parameters,
154            )
155        else:
156            log_error(
157                f"Game '{game}' is missing a StateTracker mapping in the registry.",
158                _project_parameters,
159            )
160    if game not in AVAILABLE_EMULATORS:
161        if _project_parameters["debug_mode"]:
162            log_warn(
163                f"Warning: Game '{game}' is missing an Emulator mapping in the registry.",
164                _project_parameters,
165            )
166        else:
167            log_error(
168                f"Game '{game}' is missing an Emulator mapping in the registry.",
169                _project_parameters,
170            )
171
172
173def infer_game(game: str, parameters: dict = None) -> str:
174    """
175    Try to infer the proper string identifier for a game given a possibly similar user input
176
177    Example Usage:
178    ```python
179    inferred_game = infer_game("pokemon red", parameters)
180    print(inferred_game)  # Output: "pokemon_red"
181    ```
182    Args:
183        game (str): The game variant name to infer.
184        parameters (dict): Additional parameters for logging.
185
186    Returns:
187        str: The inferred variant name.
188    """
189    parameters = load_parameters(parameters)
190    game = game.strip().lower()
191    game = game.replace(" ", "_").replace("-", "_")
192    if game in AVAILABLE_GAMES:
193        return game
194    else:
195        log_error(
196            f"Could not infer game from '{game}'. Available games are: {AVAILABLE_GAMES}",
197            parameters,
198        )
199
200
201def get_state_parser_class(
202    game: str, parameters: Optional[dict] = None
203) -> Type[StateParser]:
204    """
205    Factory method to get the strongest available StateParser class for a given game.
206
207    Args:
208        game (str): The game variant name (e.g., `pokemon_red`).
209        parameters (dict, optional): Additional parameters for logging.
210    Returns:
211        Type[StateParser]: The StateParser class for the specified game.
212    """
213    parameters = load_parameters(parameters)
214    game = infer_game(game, parameters=parameters)
215    state_parser_class = _STRONGEST_PARSERS.get(game, None)
216    if state_parser_class is None:
217        log_error(
218            f"There is no StateParser for game '{game}' in the registry.", parameters
219        )
220    return state_parser_class
221
222
223def get_state_tracker_class(
224    game: str,
225    tracker_variant: Union[str, Type[StateTracker]] = "default",
226    parameters: Optional[dict] = None,
227) -> Type[StateTracker]:
228    """
229    Factory method to get a StateTracker class for a given game and tracker variant.
230    Args:
231        game (str): The game variant name (e.g., `pokemon_red`).
232        tracker_variant (Union[str, Type[StateTracker]]): The variant of the state tracker to use. Can either be a StateTracker class (in which case it is returned directly), or a string identifier for the tracker variant (e.g., `default`).
233        parameters (dict, optional): Additional parameters for logging.
234
235    Returns:
236        Type[StateTracker]: The StateTracker class for the specified game and variant.
237    """
238    parameters = load_parameters(parameters)
239    game = infer_game(game, parameters=parameters)
240    available_trackers = AVAILABLE_STATE_TRACKERS.get(game, None)
241    if available_trackers is None:
242        log_error(
243            f"There are no available StateTrackers for game '{game}' in the registry.",
244            parameters,
245        )
246    if isinstance(tracker_variant, str):
247        if tracker_variant not in available_trackers:
248            log_error(
249                f"StateTracker variant '{tracker_variant}' is not available for game '{game}'. Available variants are: {list(available_trackers.keys())}",
250                parameters,
251            )
252        return available_trackers[tracker_variant]
253    elif issubclass(tracker_variant, StateTracker):
254        # just verify that the tracker is available for this game
255        if tracker_variant not in available_trackers.values():
256            log_error(
257                f"StateTracker class '{tracker_variant.__name__}' is not registered as an allowed tracker for game '{game}'. Available variants are: {list(available_trackers.keys())}",
258                parameters,
259            )
260        return tracker_variant
261    else:
262        log_error(
263            f"tracker_variant must either be a string identifier or a StateTracker class. Got '{type(tracker_variant)}' instead.",
264            parameters,
265        )
266
267
268def get_emulator_class(
269    game: str,
270    emulator_variant: Union[str, Type[Emulator]] = "default",
271    parameters: Optional[dict] = None,
272) -> Type[Emulator]:
273    """
274    Factory method to get an Emulator class for a given game and emulator variant.
275
276    Args:
277        game (str): The game variant name (e.g., `pokemon_red`).
278        emulator_variant (Union[str, Type[Emulator]]): The variant of the emulator to use. Can either be an Emulator class (in which case it is returned directly), or a string identifier for the emulator variant (e.g., `default`).
279        parameters (dict, optional): Additional parameters for logging.
280
281    Returns:
282        Type[Emulator]: The Emulator class for the specified game and variant.
283    """
284    parameters = load_parameters(parameters)
285    game = infer_game(game, parameters=parameters)
286    available_emulators = AVAILABLE_EMULATORS.get(game, None)
287    if available_emulators is None:
288        log_error(
289            f"There are no available Emulators for game '{game}' in the registry.",
290            parameters,
291        )
292    if isinstance(emulator_variant, str):
293        if emulator_variant not in available_emulators:
294            log_error(
295                f"Emulator variant '{emulator_variant}' is not available for game '{game}'. Available variants are: {list(available_emulators.keys())}",
296                parameters,
297            )
298        return available_emulators[emulator_variant]
299    elif issubclass(emulator_variant, Emulator):
300        # just verify that the emulator is available for this game
301        if emulator_variant not in available_emulators.values():
302            log_error(
303                f"Emulator class '{emulator_variant.__name__}' is not registered as an allowed emulator for game '{game}'. Available variants are: {list(available_emulators.keys())}",
304                parameters,
305            )
306        return emulator_variant
307    else:
308        log_error(
309            f"emulator_variant must either be a string identifier or an Emulator class. Got '{type(emulator_variant)}' instead.",
310            parameters,
311        )
312
313
314def get_emulator(
315    game: str,
316    *,
317    parameters: Optional[dict] = None,
318    init_state: str = None,
319    state_tracker_class: Union[str, Type[StateTracker]] = "default",
320    **emulator_kwargs,
321) -> Emulator:
322    """
323    Factory method to get a Pokemon emulator instance based on the specified variant.
324    Args:
325        game (str): The variant of the Pokemon game (e.g., `pokemon_red`, `pokemon_crystal`).
326        parameters (dict, optional): Additional parameters for emulator configuration.
327        init_state_name (str, optional): Name of the initial state file to load (not the path).
328        state_tracker_class (Union[str, Type[StateTracker]]): The string identifier variant of the state tracker to use, or the class itself.
329        **emulator_kwargs: Additional keyword arguments to pass to the `Emulator` constructor (e.g. `headless`)
330    Returns:
331        Emulator: An instance of the Emulator class configured for the specified variant.
332    """
333    parameters = load_parameters(parameters)
334    game = infer_game(game, parameters=parameters)
335    if f"{game}_rom_data_path" not in parameters:
336        log_error(
337            f"ROM data path for game '{game}' is not specified in the parameters under key '{game}_rom_data_path'.",
338            parameters,
339        )
340    gb_path = parameters[f"{game}_rom_data_path"] + "/" + GAME_TO_GB_NAME[game]
341    if init_state is not None:
342        if not init_state.endswith(".state"):
343            init_state = init_state + ".state"
344        init_state = parameters[f"{game}_rom_data_path"] + "/states/" + init_state
345    else:
346        init_state = parameters[f"{game}_rom_data_path"] + "/states/default.state"
347    state_parser_class = get_state_parser_class(game, parameters=parameters)
348    state_tracker_class: Type[StateTracker] = get_state_tracker_class(
349        game, tracker_variant=state_tracker_class, parameters=parameters
350    )
351    emulator_class = get_emulator_class(game, parameters=parameters)
352    emulator = emulator_class(
353        game=game,
354        gb_path=gb_path,
355        init_state=init_state,
356        state_parser_class=state_parser_class,
357        state_tracker_class=state_tracker_class,
358        parameters=parameters,
359        **emulator_kwargs,
360    )
361    return emulator
362
363
364def get_available_init_states(game: str, parameters: Optional[dict] = None) -> list:
365    """
366    Returns a list of available initial state names for the specified game.
367
368    Args:
369        game (str): The variant of the Pokemon game (e.g., `pokemon_red`, `pokemon_crystal`).
370        parameters (dict, optional): Additional parameters for configuration.
371
372    Returns:
373        list: A list of available initial state names (without .state extension).
374    """
375    parameters = load_parameters(parameters)
376    game = infer_game(game, parameters=parameters)
377    if f"{game}_rom_data_path" not in parameters:
378        log_error(
379            f"ROM data path for game '{game}' is not specified in the parameters under key '{game}_rom_data_path'.",
380            parameters,
381        )
382    states_dir = parameters[f"{game}_rom_data_path"] + "/states/"
383    if not os.path.exists(states_dir):
384        log_error(
385            f"States directory '{states_dir}' does not exist for game '{game}'.",
386            parameters,
387        )
388    state_names = [
389        f.replace(".state", "") for f in os.listdir(states_dir) if f.endswith(".state")
390    ]
391    return state_names
392
393
394def get_train_init_states(game: str, parameters: Optional[dict] = None) -> list:
395    """
396    Returns a list of allowed initial states for training agents to play the specified game.
397    This is determined based on the benchmark tasks specified for the game - any state that is a test state for a benchmark task is disallowed as a training initial state.
398
399    Args:
400        game (str): The variant of the Pokemon game (e.g., `pokemon_red`, `pokemon_crystal`).
401        parameters (dict, optional): Additional parameters for configuration.
402
403    Returns:
404        list: A list of available initial state names (without .state extension) that can be used for training.
405    """
406    parameters = load_parameters(parameters)
407    benchmark_tasks_df = get_benchmark_tasks(game, parameters=parameters)
408    test_init_states = benchmark_tasks_df["init_state"].unique().tolist()
409    other_disallowed_states = []
410    for i, row in benchmark_tasks_df.iterrows():
411        others = row["other_disallowed_states"]
412        if others and isinstance(others, str):
413            other_disallowed_states.extend(others.split(","))
414    test_init_states.extend(other_disallowed_states)
415    test_init_states = list(set(test_init_states))
416    available_init_states = get_available_init_states(game, parameters=parameters)
417    train_init_states = []
418    for state in available_init_states:
419        if (
420            state in test_init_states
421            or state.startswith("test_")
422            or "_test_" in state
423            or state.endswith("_test")
424            or state == "test"
425        ):
426            continue
427        train_init_states.append(state)
428    if len(train_init_states) == 0:
429        if parameters["debug_mode"]:
430            log_warn(
431                f"No available training initial states found for game '{game}' after filtering out test states. Returning all available initial states for now, but you should add some training states that are not used as test states or other_disallowed_states in the benchmark tasks.",
432                parameters,
433            )
434            return available_init_states
435        else:
436            log_error(
437                f"No available training initial states found for game '{game}' after filtering out test states. Please ensure that there are some initial states available for training that are not used as test states or other_disallowed_states in the benchmark tasks.",
438                parameters,
439            )
440    return train_init_states