gameboy_worlds.utils

  1# These are all the utils functions or classes that you may want to import in your project
  2from gameboy_worlds.utils.parameter_handling import load_parameters
  3from gameboy_worlds.utils.log_handling import log_error, log_info, log_warn, log_dict
  4from gameboy_worlds.utils.fundamental import file_makedir, check_optional_installs
  5from pandas import isna
  6from typing import Type, List
  7import numpy as np
  8import os
  9import matplotlib.pyplot as plt
 10import pandas as pd
 11from time import perf_counter_ns
 12from typing import Optional, List, Dict
 13
 14
 15def import_cv2(parameters: dict = None):
 16    """
 17    Import cv2 lazily so headless runs without video writing do not load extra SDL libraries.
 18    """
 19    try:
 20        import cv2
 21    except ImportError:
 22        log_error(
 23            "OpenCV (cv2) is required for emulator video recording utilities.",
 24            parameters,
 25        )
 26    return cv2
 27
 28
 29def import_pygame(parameters: dict):
 30    """
 31    Import pygame lazily so headless runs don't pull in SDL unless rendering is requested.
 32    """
 33    try:
 34        import pygame
 35    except ImportError:
 36        log_error(
 37            "pygame is required for render() / human play display. Install pygame or run without render.",
 38            parameters,
 39        )
 40    return pygame
 41
 42
 43def is_none_str(s) -> bool:
 44    """
 45    Checks if a string is None or represents a null value.
 46
 47    Args:
 48        s (str or None): The string to check.
 49
 50    Returns:
 51        bool: True if the string is None or represents a null value, False otherwise.
 52    """
 53    if s is None:
 54        return True
 55    if isinstance(s, str):
 56        options = ["none", "null", "nan", ""]
 57        for option in options:
 58            if s.lower() == option:
 59                return True
 60    return isna(s)
 61
 62
 63def nested_dict_to_str(
 64    nested_dict: dict, *, indent: int = 0, indent_char: str = "  "
 65) -> str:
 66    """
 67    Converts a nested dictionary to a formatted string representation.
 68    Example Usage:
 69    ```python
 70    nested_dict={2: 4, 3: {4: 5, 6: {7: 8}}}
 71    print(nested_dict_to_str(nested_dict))
 72    2: 4
 73    3: Dict:
 74      4: 5
 75      6: Dict:
 76        7: 8
 77    ```
 78
 79    Args:
 80        nested_dict (dict): The nested dictionary to convert.
 81        indent (int): The current indentation level.
 82        indent_char (str): The character(s) used for indentation.
 83    Returns:
 84        str: A formatted string representation of the nested dictionary.
 85
 86    """
 87    result = ""
 88    for key, value in nested_dict.items():
 89        result += indent_char * indent + str(key) + ": "
 90        if isinstance(value, dict):
 91            result += "Dict: \n" + nested_dict_to_str(
 92                value, indent + 1, indent_char=indent_char
 93            )
 94        else:
 95            result += str(value) + "\n"
 96    return result
 97
 98
 99def verify_parameters(parameters: dict):
100    """
101    Does a basic sanity check to ensure parameters is a non-empty dictionary.
102    """
103    if parameters is None:
104        raise ValueError("Parameters cannot be None.")
105    if not isinstance(parameters, dict):
106        raise ValueError("Parameters must be a dictionary.")
107    if len(parameters) == 0:
108        raise ValueError("Parameters dictionary cannot be empty.")
109
110
111def get_lowest_level_subclass(class_list: List[Type]) -> Type:
112    """
113    Given a list of classes, returns the class that is the lowest level subclass in the inheritance hierarchy.
114    """
115    lowest_level_tracker = None
116    for cls in class_list:
117        if lowest_level_tracker is None:
118            lowest_level_tracker = cls
119        elif issubclass(cls, lowest_level_tracker):
120            lowest_level_tracker = cls
121    return lowest_level_tracker
122
123
124def show_frames(
125    frames: np.ndarray, titles: List[str] = None, save=False, parameters: dict = None
126):
127    """
128    Plots each frame as an image in matplotlib. If save is true, will save each frame as title.png in the frame_saves/ directory.
129    titles length must be equal to frame length if specified.
130    """
131    parameters = load_parameters(parameters)
132    if isinstance(frames, list):
133        for i in range(len(frames)):
134            if frames[i].ndim == 2:
135                frames[i] = np.expand_dims(frames[i], axis=-1)
136    else:
137        if not isinstance(frames, np.ndarray):
138            log_error(
139                f"Frames must be a numpy array or list of numpy arrays, but got {type(frames)}",
140                parameters,
141            )
142        if frames.ndim == 2:
143            frames = [np.expand_dims(frames, axis=-1)]
144        elif frames.ndim == 3:
145            # now either we have (num_frames, height, width) or (height, width, channels)
146            if frames.shape[2] == 1:
147                frames = [frames]
148            else:
149                frames = [
150                    np.expand_dims(frames[i], axis=-1) for i in range(frames.shape[0])
151                ]
152    if isinstance(titles, str):
153        titles = [titles]
154    if save:
155        if titles is None:
156            log_error(f"Cannot save frames without titles specified.", parameters)
157    if titles is not None:
158        if len(titles) == 1 and len(frames) > 1:
159            titles = [titles[0] + f"_{i}" for i in range(len(frames))]
160    if titles is not None and len(titles) != len(frames):
161        log_error(
162            f"Length of titles {len(titles)} does not match number of frames {len(frames)}",
163            parameters,
164        )
165    save_dir = "frame_saves/"
166    os.makedirs(save_dir, exist_ok=True)
167
168    for i in range(len(frames)):
169        plt.imshow(frames[i])
170        if titles is not None:
171            plt.title(titles[i])
172        if save:
173            filename = os.path.join(
174                save_dir, titles[i].replace(" ", "_").replace("/", "_") + ".png"
175            )
176            plt.imsave(filename, frames[i][:, :, 0], cmap="gray")
177        else:
178            plt.show()
179
180
181def get_train_games(game: str, parameters: dict = None) -> List[str]:
182    """
183    Returns a list of games that can be used for training the specified game variant.
184
185    Args:
186        game (str): The variant of the game to get training games for.
187        parameters (dict, optional): Additional parameters for error logging.
188    """
189    parameters = load_parameters(parameters)
190    tasks_df = get_benchmark_tasks(game, parameters, shifted_included=True)
191    training_rows = tasks_df[tasks_df["can_train_from_init_state"] == True]
192    if len(training_rows) == 0:
193        log_error(
194            f"No training games specified for game variant '{game}' in benchmark tasks. This is an error.",
195            parameters,
196        )
197    training_games = list(set(training_rows["game"].tolist()))
198    return training_games
199
200
201def get_benchmark_tasks_dfs(parameters: dict = None) -> dict[str, pd.DataFrame]:
202    """
203    Loads the benchmark tasks from the benchmark/tests/tasks.csv file
204
205    Args:
206        parameters (dict, optional): Additional parameters for error logging.
207    Returns:
208        dict[str, pd.DataFrame]: A dictionary mapping game names to their corresponding DataFrames containing the benchmark tasks.
209    """
210    parameters = load_parameters(parameters)
211    project_root = parameters["project_root"]
212    module_paths = os.listdir(project_root + "/benchmark/tests/")
213    benchmark_dfs = {}
214    for module_path in module_paths:
215        if module_path.endswith(".csv"):
216            tasks_filepath = os.path.join(project_root, "benchmark/tests", module_path)
217            benchmark_name = module_path.strip(".csv")
218            benchmark_dfs[benchmark_name] = pd.read_csv(tasks_filepath)
219    return benchmark_dfs
220
221
222def get_benchmark_tasks(
223    game: str, parameters: dict = None, shifted_included: bool = False
224) -> pd.DataFrame:
225    """
226    Loads the benchmark tasks for the specified game from the benchmark/tests/tasks.csv file
227
228    Args:
229        game (str): The variant of the game to get benchmark tasks for.
230        parameters (dict, optional): Additional parameters for error logging.
231        shifted_included (bool, optional): Whether to include task rows from other titles in the same series as part of the benchmark tasks.
232
233    Returns:
234        pd.DataFrame: DataFrame containing the benchmark tasks for the specified game.
235    """
236    parameters = load_parameters(parameters)
237    tasks_dfs = get_benchmark_tasks_dfs(parameters)
238    available_games = set()
239    task_df = None
240    for benchmark_name, df in tasks_dfs.items():
241        available_games.update(set(df["game"].unique()))
242        if game in df["game"].unique():
243            if task_df is not None:
244                log_error(
245                    f"Multiple benchmark modules contain tasks for game variant '{game}'. Please ensure that only one benchmark CSV file contains tasks for this game.",
246                    parameters,
247                )
248            task_df = df
249
250    if game not in available_games:
251        log_error(
252            f"Game variant '{game}' not found in benchmark tasks. Available game variants: {available_games}",
253            parameters=parameters,
254        )
255    if not shifted_included:
256        game_tasks_df = task_df[task_df["game"] == game].reset_index(drop=True)
257    else:
258        game_tasks_df = task_df
259    return game_tasks_df
260
261
262def get_training_states(game: str, parameters: dict = None) -> Optional[List[str]]:
263    """
264    Loads the training states for the specified game from the benchmark/tasks.csv file
265
266    Args:
267        game (str): The variant of the game to get training states for.
268        parameters (dict, optional): Additional parameters for error logging.
269
270    Returns:
271        Optional[List[str]]: A list of training states for the specified game, or None if no training states are specified.
272    """
273    parameters = load_parameters(parameters)
274    tasks_df = get_benchmark_tasks(game, parameters, shifted_included=False)
275    training_rows = tasks_df[tasks_df["can_train_from_init_state"] == True]
276    if len(training_rows) == 0:
277        log_warn(
278            f"No training states specified for game variant '{game}' in benchmark tasks.",
279            parameters,
280        )
281        return None
282    training_states = list(set(training_rows["init_state"].tolist()))
283    return training_states
284
285
286def get_shifted_training_states(
287    game: str, parameters: dict = None
288) -> Dict[str, List[str]]:
289    """
290    Loads the shifted training states for the specified game from the benchmark/tasks.csv file.
291
292    This ends up being all available training states from every other title of the same series. So if you ask for pokemon_red's shifted states, you will get pokemon_crystal's training states.
293
294    Args:
295        game (str): The variant of the game to get shifted training states for.
296        parameters (dict, optional): Additional parameters for error logging.
297    Returns:
298        Dict[str, List[str]]: A dictionary of lists of shifted training states for the specified game. Each entry is in format {game: [training_states]}, where game is a different title in the same series as the input game.
299    """
300    parameters = load_parameters(parameters)
301    tasks_df = get_benchmark_tasks(game, parameters, shifted_included=True)
302    other_task_rows = tasks_df[tasks_df["game"] != game].reset_index(drop=True)
303    training_rows = other_task_rows[
304        other_task_rows["can_train_from_init_state"] == True
305    ]
306    if len(training_rows) == 0:
307        log_error(
308            f"No shifted training states specified for game variant '{game}' in benchmark tasks. This shouldn't happen.",
309            parameters,
310        )
311    shifted_training_states = {}
312    for _, row in training_rows.iterrows():
313        other_game = row["game"]
314        if other_game not in shifted_training_states:
315            shifted_training_states[other_game] = []
316        shifted_training_states[other_game].append(row["init_state"])
317    return shifted_training_states
318
319
320def get_all_training_states(parameters: dict = None) -> Dict[str, List[str]]:
321    """
322    Gets all regular training states for the all games (for which training states are specified).
323
324    Args:
325        parameters (dict, optional): Additional parameters for error logging.
326
327    Returns:
328        Dict[str, List[str]]: A dictionary containing {game: training_states} entries for all games for which training states are specified in the benchmark tasks.
329    """
330    parameters = load_parameters(parameters)
331    tasks_dfs = get_benchmark_tasks_dfs(parameters)
332    all_games = set()
333    for df in tasks_dfs.values():
334        all_games.update(df["game"].unique())
335    all_training_states = {}
336    for game in all_games:
337        training_states = get_training_states(game, parameters)
338        if training_states is not None:
339            all_training_states[game] = training_states
340    return all_training_states
341
342
343def get_all_shifted_training_states(
344    parameters: dict = None,
345) -> Dict[str, Dict[str, List[str]]]:
346    """
347    Gets all shifted training states for all games.
348
349    Args:
350        parameters (dict, optional): Additional parameters for error logging.
351
352    Returns:
353        Dict[str, Dict[str, List[str]]]: A nested dictionary containing {game: {other_game: shifted_training_states}} entries for all games for which shifted training states are specified in the benchmark tasks.
354    """
355    parameters = load_parameters(parameters)
356    tasks_dfs = get_benchmark_tasks_dfs(parameters)
357    all_games = set()
358    for df in tasks_dfs.values():
359        all_games.update(df["game"].unique())
360    all_shifted_training_states = {}
361    for game in all_games:
362        shifted_states = get_shifted_training_states(game, parameters)
363        all_shifted_training_states[game] = shifted_states
364    return all_shifted_training_states
365
366
367class _Profiler:
368    """
369    A simple profiler class to track the time taken by different events in the code. It can also group events together and show the percentage of time taken by each event in the group.
370    """
371
372    LOG_EVENTS = False
373    last_event_time = None
374    last_event = None
375    group = None
376    group_name = None
377
378    @staticmethod
379    def event(event_name):
380        current_time = perf_counter_ns()
381        if _Profiler.last_event_time is not None:
382            elapsed_time = current_time - _Profiler.last_event_time
383            if _Profiler.LOG_EVENTS:
384                log_info(
385                    f"{_Profiler.last_event} -> {event_name}: {elapsed_time / 1e6:.2f} ms",
386                )
387        _Profiler.last_event_time = current_time
388        _Profiler.last_event = event_name
389        if _Profiler.group_name is not None:
390            _Profiler.group.append((event_name, current_time))
391
392    @staticmethod
393    def start_group(group_name):
394        if _Profiler.group_name is not None:
395            fractions = []
396            total_time = 0
397            for i in range(1, len(_Profiler.group)):
398                event_name, event_time = _Profiler.group[i]
399                prev_event_name, prev_event_time = _Profiler.group[i - 1]
400                elapsed_time = event_time - prev_event_time
401                total_time += elapsed_time
402            for i in range(1, len(_Profiler.group)):
403                event_name, event_time = _Profiler.group[i]
404                prev_event_name, prev_event_time = _Profiler.group[i - 1]
405                elapsed_time = event_time - prev_event_time
406                fractions.append(elapsed_time / total_time if total_time > 0 else 0)
407                log_info(
408                    f"{_Profiler.group_name} - {prev_event_name} -> {event_name}: {elapsed_time / 1e6:.2f} ms ({fractions[-1]*100:.2f}%)",
409                )
410
411        _Profiler.group_name = group_name
412        _Profiler.group = []
413
414    @staticmethod
415    def close_group():
416        _Profiler.start_group(None)