gameboy_worlds.emulation.emulator

   1from enum import Enum
   2from typing import Type, Optional, Tuple
   3
   4
   5import os
   6import re
   7from time import perf_counter
   8import sys
   9import shutil
  10import uuid
  11from gameboy_worlds.emulation.parser import StateParser
  12from gameboy_worlds.emulation.tracker import StateTracker
  13from gameboy_worlds.utils import (
  14    load_parameters,
  15    log_error,
  16    log_warn,
  17    file_makedir,
  18    log_info,
  19    is_none_str,
  20    verify_parameters,
  21    log_dict,
  22    import_cv2,
  23)
  24
  25
  26from pyboy import PyBoy
  27from pyboy.utils import WindowEvent
  28from matplotlib import pyplot as plt
  29from skimage.transform import downscale_local_mean
  30import numpy as np
  31from tqdm import tqdm
  32
  33
  34class LowLevelActions(Enum):
  35    """
  36    Enum for low-level actions that can be performed on the GameBoy emulator.
  37    """
  38
  39    PRESS_ARROW_DOWN = WindowEvent.PRESS_ARROW_DOWN
  40    PRESS_ARROW_LEFT = WindowEvent.PRESS_ARROW_LEFT
  41    PRESS_ARROW_RIGHT = WindowEvent.PRESS_ARROW_RIGHT
  42    PRESS_ARROW_UP = WindowEvent.PRESS_ARROW_UP
  43    PRESS_BUTTON_A = WindowEvent.PRESS_BUTTON_A
  44    PRESS_BUTTON_B = WindowEvent.PRESS_BUTTON_B
  45    PRESS_BUTTON_START = WindowEvent.PRESS_BUTTON_START
  46    # PRESS_BUTTON_SELECT = WindowEvent.PRESS_BUTTON_SELECT
  47
  48
  49class ReleaseActions(Enum):
  50    """
  51    Enum for release actions corresponding to low-level actions.
  52    """
  53
  54    release_actions = {
  55        LowLevelActions.PRESS_ARROW_DOWN: WindowEvent.RELEASE_ARROW_DOWN,
  56        LowLevelActions.PRESS_ARROW_LEFT: WindowEvent.RELEASE_ARROW_LEFT,
  57        LowLevelActions.PRESS_ARROW_RIGHT: WindowEvent.RELEASE_ARROW_RIGHT,
  58        LowLevelActions.PRESS_ARROW_UP: WindowEvent.RELEASE_ARROW_UP,
  59        LowLevelActions.PRESS_BUTTON_A: WindowEvent.RELEASE_BUTTON_A,
  60        LowLevelActions.PRESS_BUTTON_B: WindowEvent.RELEASE_BUTTON_B,
  61        LowLevelActions.PRESS_BUTTON_START: WindowEvent.RELEASE_BUTTON_START,
  62        # LowLevelActions.PRESS_BUTTON_SELECT: WindowEvent.RELEASE_BUTTON_SELECT,
  63    }
  64
  65
  66class IDPathCreator:
  67    """
  68    Handles the creation of IDs and paths for saving emulator artifacts.
  69    """
  70
  71    def __init__(self, parameters: dict):
  72        verify_parameters(parameters)
  73        self._parameters = parameters
  74
  75    def _get_numbered_instance_id(self, path: str, instance_id: str = None) -> int:
  76        """
  77        Looks at the given path and returns <next_number>_<instance_id> where next_number is 1 + the highest existing numbered instance ID in the path.
  78
  79        :param path: Path to look for existing instances.
  80        :type path: str
  81        :param instance_id: Instance ID pattern to match. If None, counts all instances.
  82        :type instance_id: str
  83        :return: Number of existing instances matching the pattern.
  84        """
  85        # instance_id = str(uuid.uuid4())[:8]
  86        if not os.path.exists(path):
  87            if instance_id is None:
  88                instance_id = str(uuid.uuid4())[:8]
  89            return f"0_{instance_id}"
  90        if instance_id is None:
  91            instance_id = str(uuid.uuid4())[:8]
  92            n_existing = os.listdir(path)
  93            return f"{len(n_existing)}_{instance_id}"
  94        # must find the pattern <number>_<instance_id>
  95        pattern = re.compile(r"(\d+)_" + re.escape(instance_id))
  96        existing_instances = [d for d in os.listdir(path) if pattern.match(d)]
  97        if len(existing_instances) == 0:
  98            return f"0_{instance_id}"
  99        else:
 100            return f"{len(existing_instances)}_{instance_id}"
 101
 102    def get_session_path(
 103        self,
 104        session_name: Optional[str],
 105        instance_id: Optional[str],
 106        environment_variant: str,
 107    ) -> str:
 108        if session_name is None:
 109            session_name = "tmp_sessions"
 110            log_warn(
 111                f"Saving a temporary session. If you run gameboy_worlds.clear_tmp_sessions(), it will be deleted. To make it permanent, pass in a `session_name` to the emulator constructor kwargs."
 112            )
 113        elif not isinstance(session_name, str) or session_name == "":
 114            log_error(
 115                f"session_name must be a non-empty string. Recieved {session_name}",
 116                self._parameters,
 117            )
 118        storage_dir = self._parameters["storage_dir"]
 119        session_path = os.path.join(
 120            storage_dir, "sessions", environment_variant, session_name
 121        )
 122        if instance_id is not None:
 123            if not isinstance(instance_id, str) or instance_id == "":
 124                log_error(
 125                    f"instance_id must be a non-empty string. Recieved {instance_id}",
 126                    self._parameters,
 127                )
 128            session_path = os.path.join(session_path, "named_instances")
 129        instance_id = self._get_numbered_instance_id(session_path, instance_id)
 130        full_session_path = os.path.join(session_path, instance_id)
 131        os.makedirs(full_session_path, exist_ok=True)
 132        return full_session_path
 133
 134    def clear_tmp_sessions(self):
 135        """
 136        Clears the tmp_sessions directory for ALL game variants.
 137        """
 138        storage_dir = self._parameters["storage_dir"]
 139        session_path = os.path.join(storage_dir, "sessions")
 140        if not os.path.exists(session_path):
 141            return
 142        existing_variants = os.listdir(session_path)
 143        cleared_sessions = {}
 144        for variant in existing_variants:
 145            tmp_sessions_path = os.path.join(session_path, variant, "tmp_sessions")
 146            if not os.path.exists(tmp_sessions_path):
 147                continue
 148            n_sessions = os.listdir(tmp_sessions_path)
 149            shutil.rmtree(tmp_sessions_path)
 150            cleared_sessions[variant] = len(n_sessions)
 151        log_info(
 152            f"Cleared temporary sessions, statistics:",
 153            self._parameters,
 154        )
 155        log_dict(cleared_sessions, parameters=self._parameters)
 156
 157
 158class VideoWriter:
 159    def __init__(
 160        self,
 161        *,
 162        session_path: str,
 163        output_shape: Tuple[int, int],
 164        reduce_resolution: bool,
 165        parameters: dict,
 166    ):
 167        verify_parameters(parameters)
 168        self._session_path = session_path
 169        self._output_shape = output_shape
 170        self._reduce_resolution = reduce_resolution
 171        self._parameters = parameters
 172        self._frame_writer = None
 173        self.video_running = False
 174        """ Whether the video writer is currently recording video. """
 175        project_dir = self._parameters["project_root"]
 176        cv2 = import_cv2(self._parameters)
 177        self._button_images = {
 178            None: cv2.imread(
 179                os.path.join(project_dir, "assets/buttons/idle.png"),
 180                cv2.IMREAD_UNCHANGED,
 181            ),
 182            LowLevelActions.PRESS_ARROW_DOWN: cv2.imread(
 183                os.path.join(project_dir, "assets/buttons/down.png"),
 184                cv2.IMREAD_UNCHANGED,
 185            ),
 186            LowLevelActions.PRESS_ARROW_LEFT: cv2.imread(
 187                os.path.join(project_dir, "assets/buttons/left.png"),
 188                cv2.IMREAD_UNCHANGED,
 189            ),
 190            LowLevelActions.PRESS_ARROW_RIGHT: cv2.imread(
 191                os.path.join(project_dir, "assets/buttons/right.png"),
 192                cv2.IMREAD_UNCHANGED,
 193            ),
 194            LowLevelActions.PRESS_ARROW_UP: cv2.imread(
 195                os.path.join(project_dir, "assets/buttons/up.png"), cv2.IMREAD_UNCHANGED
 196            ),
 197            LowLevelActions.PRESS_BUTTON_A: cv2.imread(
 198                os.path.join(project_dir, "assets/buttons/a.png"), cv2.IMREAD_UNCHANGED
 199            ),
 200            LowLevelActions.PRESS_BUTTON_B: cv2.imread(
 201                os.path.join(project_dir, "assets/buttons/b.png"), cv2.IMREAD_UNCHANGED
 202            ),
 203            LowLevelActions.PRESS_BUTTON_START: cv2.imread(
 204                os.path.join(project_dir, "assets/buttons/start.png"),
 205                cv2.IMREAD_UNCHANGED,
 206            ),
 207            # LowLevelActions.PRESS_BUTTON_SELECT: cv2.imread(
 208            #     os.path.join(project_dir, "assets/buttons/select.png"), cv2.IMREAD_UNCHANGED
 209            # ),
 210        }
 211
 212    def _get_free_video_id(self) -> str:
 213        """
 214        Returns a new unique video ID for saving video files.
 215
 216        Returns:
 217            str: A new unique video ID.
 218        """
 219        base_dir = os.path.join(self._session_path, "videos")
 220        videos = os.listdir(base_dir) if os.path.exists(base_dir) else []
 221        # all will be something.mp4, if its int.mp4, get the int
 222        video_ints = []
 223        for video in videos:
 224            if video.endswith(".mp4"):
 225                video_name = video[:-4]
 226                if video_name.isdigit():
 227                    video_ints.append(int(video_name))
 228        if len(video_ints) == 0:
 229            return "0.mp4"
 230        return str(max(video_ints) + 1) + ".mp4"
 231
 232    def start_video(self, video_id: str = None):
 233        """
 234        Starts recording video of the emulator's screen.
 235        Args:
 236            video_id (str, optional): Name of the video file to save. If None, a new unique name will be generated.
 237        """
 238        if video_id is not None:
 239            if not isinstance(video_id, str):
 240                log_error(
 241                    "video_id must be a string (not digits) if provided.",
 242                    self._parameters,
 243                )
 244            if not video_id.endswith(".mp4"):
 245                log_error("video_id must end with .mp4 if provided.", self._parameters)
 246            if os.path.exists(os.path.join(self._session_path, "videos", video_id)):
 247                log_warn(
 248                    f"video_id {video_id} already exists. Overwriting...",
 249                    self._parameters,
 250                )
 251        else:
 252            video_id = self._get_free_video_id()
 253        base_dir = os.path.join(self._session_path, "videos")
 254        os.makedirs(base_dir, exist_ok=True)
 255        video_path = os.path.join(base_dir, f"{video_id}")
 256        self.close_video()
 257        cv2 = import_cv2(self._parameters)
 258        self._frame_writer = cv2.VideoWriter(
 259            video_path,
 260            cv2.VideoWriter_fourcc(*"mp4v"),
 261            60,
 262            (self._output_shape[0], self._output_shape[1]),
 263            isColor=True,
 264        )
 265        self.video_running = True
 266        log_info(f"\nStarted recording video to: {video_path}\n", self._parameters)
 267
 268    def _get_reduced(self, frame: np.ndarray) -> np.ndarray:
 269        """
 270        Reduces the resolution of the given frame by a factor of 2 using local mean downscaling.
 271        Args:
 272            frame (np.ndarray): The frame to reduce the resolution of.
 273        Returns:
 274            np.ndarray: The reduced resolution frame.
 275        """
 276        reduced = (downscale_local_mean(frame, (2, 2, 1))).astype(np.uint8)
 277        return reduced
 278
 279    def add_video_frames(
 280        self, frames: np.ndarray, pressed_button: Optional[LowLevelActions] = None, show_button: bool = True
 281    ):
 282        """
 283        Adds a list of frame from the emulator to the video being recorded.
 284
 285        Args:
 286            frames (np.ndarray): A stack of frames to add to the video. Shape is [n_frames, height, width, channels].
 287            pressed_button (LowLevelActions, optional): The button that was pressed during these frames. If None, no button overlay will be added.
 288            show_button (bool, optional): Whether to show the button overlay on the video. Defaults to True.
 289        """
 290        if show_button:
 291            button_image = self._button_images[pressed_button]
 292            button_size = 50
 293            button_offset = 0
 294            button_x = self._output_shape[0] - button_size - button_offset
 295            button_y = self._output_shape[1] - button_size - button_offset
 296            cv2 = import_cv2(self._parameters)
 297            button_image = cv2.resize(button_image, (button_size, button_size))
 298            alphas = button_image[:, :, 3] / 255.0
 299
 300        for frame in frames:
 301            if self._reduce_resolution:
 302                frame = self._get_reduced(frame)
 303            # frame_size = (current_frame.shape[1], current_frame.shape[0], 1) # Width, Height, should be equal to self.output_shape
 304            # expand grayscale frame to 3 channels for video writing
 305            treated_frame = np.repeat(frame, 3, axis=2)
 306            if show_button:
 307                treated_frame[
 308                    button_y : button_y + button_size, button_x : button_x + button_size
 309                ] = (
 310                    treated_frame[
 311                        button_y : button_y + button_size, button_x : button_x + button_size
 312                    ]
 313                    * (1 - alphas[:, :, np.newaxis])
 314                    + button_image[:, :, :3] * alphas[:, :, np.newaxis]
 315                ).astype(
 316                    np.uint8
 317                )
 318            else:
 319                pass
 320            self._frame_writer.write(treated_frame)
 321        return
 322
 323    def close_video(self):
 324        """
 325        Closes the video writer and stops recording video.
 326        """
 327        if self._frame_writer is not None:
 328            self._frame_writer.release()
 329            self._frame_writer = None
 330        self.video_running = False
 331
 332
 333class Emulator:
 334    """
 335    Handles the running of the GameBoy emulator, including loading ROMs, managing state, performing low level actions and calling the tracker.
 336    Subclasses will likely only be needed to manually force the execution of specific button sequences when certain states are detected. (e.g. short-circuiting specific menus, etc.)
 337
 338    Can be used to access the `state_parser` and `state_tracker` instances for the running game instance.
 339    """
 340
 341    REQUIRED_STATE_PARSER = StateParser
 342    """ The minimal functionality StateParser needed for this emulator to run """
 343
 344    REQUIRED_STATE_TRACKER = StateTracker
 345    """ The minimal functionality StateTracker needed for this emulator to run """
 346
 347    def __init__(
 348        self,
 349        game: str,
 350        gb_path: str,
 351        state_parser_class: Type[StateParser],
 352        state_tracker_class: Type[StateTracker],
 353        init_state: str,
 354        parameters: dict,
 355        *,
 356        headless: bool = True,
 357        max_steps: int = None,
 358        save_video: bool = None,
 359        session_name: str = None,
 360        instance_id: str = None,
 361        wait_ticks: int = None,
 362        press_step: int = None,
 363    ):
 364        """
 365        Start the GameBoy emulator with the given ROM file and initial state.
 366
 367        Args:
 368            game (str): Name of game variant being emulated.
 369            gb_path (str): Path to the GameBoy ROM file.
 370            state_parser_class (Type[StateParser]): A class that inherits from StateParser to parse game state variables.
 371            state_tracker_class (Type[StateTracker]): A class that inherits from StateTracker to track game state metrics.
 372            init_state (str): Path to the initial state file to load.
 373            parameters (dict): Dictionary of parameters for the environment.
 374            headless (bool, optional): Whether to run the environment in headless mode.
 375            max_steps (int, optional): Maximum number of steps per episode.
 376            save_video (bool, optional): Whether to save video of the episodes.
 377            session_name (str, optional): Name of the session. If None, a new session name will be allocated. This is the broad category you want to save files to
 378            instance_id (str, optional): Unique identifier for this environment instance. If None, a new UUID will be generated. The instance ID is useful for distinguishing multiple environments running in parallel with the same session name.
 379            wait_ticks (int, optional): Number of ticks to wait between actions.
 380            press_step (int, optional): Number of steps to press a button for.
 381        """
 382        verify_parameters(parameters)
 383        self._parameters = parameters
 384        if game is None or game == "":
 385            log_error(
 386                "You must provide a name for the game variant being emulated.",
 387                self._parameters,
 388            )
 389        if gb_path is None:
 390            log_error(
 391                "You must provide a path to the GameBoy ROM file.", self._parameters
 392            )
 393        if not issubclass(state_parser_class, self.REQUIRED_STATE_PARSER):
 394            log_error(
 395                f"state_parser_class must be a subclass of {self.REQUIRED_STATE_PARSER.__name__}, got {state_parser_class_.__name__}.",
 396                self._parameters,
 397            )
 398        if not issubclass(state_tracker_class, self.REQUIRED_STATE_TRACKER):
 399            log_error(
 400                f"state_tracker_class must be a subclass of {self.REQUIRED_STATE_TRACKER.__name__}, got {state_tracker_class.__name__}.",
 401                self._parameters,
 402            )
 403        if init_state is None:
 404            log_error(
 405                "You must provide an initial state file to load.", self._parameters
 406            )
 407        if headless not in [True, False]:
 408            log_error("headless must be a boolean.", self._parameters)
 409        self.game = game
 410        """ Name of game variant being emulated. """
 411        self._gb_path = gb_path
 412        self._set_init_state(init_state)
 413        # validate init_state exists and ends with .state
 414        if not os.path.exists(self._gb_path):
 415            log_error(
 416                f"GameBoy ROM file {self._gb_path} does not exist. You must obtain a ROM through official means, and then place it in the path: {self._gb_path}",
 417                self._parameters,
 418            )
 419        if not self._gb_path.endswith(".gb") and not self._gb_path.endswith(".gbc"):
 420            log_error(
 421                f"GameBoy ROM file {self._gb_path} is not a .gb or .gbc file.",
 422                self._parameters,
 423            )
 424        self.headless = headless
 425        """ Whether to run the environment in headless mode."""
 426        if max_steps is None:
 427            max_steps = self._parameters["gameboy_max_steps"]
 428        if max_steps > self._parameters["gameboy_hard_max_steps"]:
 429            log_warn(
 430                f"max_steps {max_steps} exceeds gameboy_hard_max_steps {self._parameters['gameboy_hard_max_steps']}. Setting to hard max.",
 431                self._parameters,
 432            )
 433            max_steps = self._parameters["gameboy_hard_max_steps"]
 434        self.max_steps = max_steps
 435        """ Maximum number of steps per episode. """
 436        id_path_creator = IDPathCreator(self._parameters)
 437        self.session_path = id_path_creator.get_session_path(
 438            session_name=session_name,
 439            instance_id=instance_id,
 440            environment_variant=self.get_env_variant(),
 441        )
 442        """ Path to the session directory. This is where all artifacts for this session are saved. """
 443
 444        if wait_ticks is None:
 445            wait_ticks = parameters["gameboy_wait_ticks"]
 446        self.wait_ticks = wait_ticks
 447        """ Number of emulator ticks to wait after an action. Defaults to value specified in config files. """
 448        if press_step is None:
 449            press_step = parameters["gameboy_press_step"]
 450        self.press_step = press_step
 451        """ Number of emulator ticks to hold down a button press. Defaults to value specified in config files. """
 452        self.render_headless = parameters["gameboy_headless_render"]
 453        """ Whether to render the emulator screen even in headless mode. This must be true for methods that rely on image observations (e.g. VLMs) to access the screen. Defaults to value specified in config files. """
 454        if not self.render_headless:
 455            log_error(
 456                "render_headless cannot be set to False. In the Pokemon environments, screen captures are used aggressively to determine state. ",
 457                self._parameters,
 458            )
 459
 460        self.reset_count = 0
 461        """ Number of times the environment has been reset. """
 462        self.step_count = 0
 463        """ Number of steps taken in the current episode. """
 464        self._reduce_video_resolution = parameters["gameboy_reduce_video_resolution"]
 465        frame_size = (
 466            160,
 467            144,
 468        )
 469        self.screen_shape = (frame_size[0], frame_size[1], 1)
 470        """ Resolution of the rendered game screen """
 471        if self._reduce_video_resolution:
 472            self.output_shape = (frame_size[0] // 2, frame_size[1] // 2)
 473        else:
 474            self.output_shape = (frame_size[0], frame_size[1])
 475            """ Shape of the output observations. This is the resolution of the rendered screen. """
 476
 477        if save_video is None:
 478            save_video = self._parameters["gameboy_default_save_video"]
 479        self.save_video = save_video
 480        """ Whether to save video of the episodes. """
 481        self.video_writer = VideoWriter(
 482            session_path=self.session_path,
 483            output_shape=self.output_shape,
 484            reduce_resolution=self._reduce_video_resolution,
 485            parameters=self._parameters,
 486        )
 487        """ Holds the VideoWriter of this Emulator instance """
 488
 489        head = "null" if self.headless else "SDL2"
 490
 491        self._pyboy = PyBoy(
 492            self._gb_path,
 493            window=head,
 494        )
 495        self.state_parser = state_parser_class(self._pyboy, self._parameters)
 496        """ Instance of the StateParser to parse game state variables. """
 497
 498        self.state_tracker = state_tracker_class(
 499            self.state_parser,
 500            self._parameters,
 501        )
 502        """ Instance of the StateTracker to track game state metrics. """
 503
 504        # self.screen = self.pyboy.botsupport_manager().screen()
 505
 506        if not self.headless:
 507            if not is_none_str(self._parameters["gameboy_headed_emulation_speed"]):
 508                self._pyboy.set_emulation_speed(
 509                    int(self._parameters["gameboy_headed_emulation_speed"])
 510                )
 511        self.reset()
 512
 513    @staticmethod
 514    def create_first_state(gb_path: str, state_path: str):
 515        """
 516        Creates a basic state for the emulator. This can be used to create an initial, default state file for a new game.
 517
 518        Warning: This method uses parameter free logging, so if you override the log_file with a command prompt argument, it will be ignored here.
 519
 520        Args:
 521            gb_path (str): Path to the GameBoy ROM file.
 522            state_path (str): Path to save the initial state file.
 523        """
 524        # error out if gb_path does not exist or is not a .gb or .gbc file
 525        if not os.path.exists(gb_path):
 526            log_error(
 527                f"GameBoy ROM file {gb_path} does not exist. You must obtain a ROM through official means, and then place it in the path: {gb_path}"
 528            )
 529        if not gb_path.endswith(".gb") and not gb_path.endswith(".gbc"):
 530            log_error(f"GameBoy ROM file {gb_path} is not a .gb or .gbc file.")
 531        if not state_path.endswith(".state"):
 532            state_path = state_path + ".state"
 533        if os.path.exists(state_path):
 534            log_error(f"State file {state_path} already exists. Will not overwrite...")
 535        file_makedir(state_path)
 536        pyboy = PyBoy(
 537            gb_path,
 538            window="null",
 539        )
 540        with open(state_path, "wb") as f:
 541            pyboy.save_state(f)
 542        pyboy.stop()
 543        log_info(f"Created initial state file at {state_path}")
 544        sys.exit(0)
 545
 546    def set_init_state(self, init_state: str):
 547        """Sets a new initial state file for the environment. and resets the environment.
 548
 549        Args:
 550            init_state (str): Path to the new initial state file.
 551        """
 552        self._set_init_state(init_state)
 553        self.reset()
 554
 555    def _set_init_state(self, init_state: str):
 556        """
 557        Sets a new initial state file for the environment to eventually load.
 558        Does not reset the environment.
 559
 560        Args:
 561            init_state (str): Path to the new initial state file.
 562        """
 563        if init_state is None:
 564            # log_warn(f"No initial state file provided. Using default initial state.", self._parameters)
 565            return
 566        if not init_state.endswith(".state"):
 567            init_state = init_state + ".state"
 568        states_path = self._parameters[f"{self.game}_rom_data_path"] + "/states/"
 569        if not os.path.exists(init_state):
 570            if states_path in init_state:
 571                log_error(
 572                    f"Initial state file {init_state} does not exist.", self._parameters
 573                )
 574            # try to resolve.
 575            potential_path = os.path.join(states_path, init_state)
 576            if os.path.exists(potential_path):
 577                init_state = potential_path
 578            else:
 579                log_error(
 580                    f"Initial state file {init_state} does not exist.", self._parameters
 581                )
 582        self.init_state = init_state
 583        log_info(f"Set new initial state file to {self.init_state}", self._parameters)
 584
 585    def reset(self, new_init_state: str = None):
 586        """
 587        Resets the environment to the initial state. Optionally loads a new initial state file.
 588
 589        Args:
 590            new_init_state (str, optional): Path to a new initial state file to load.
 591        """
 592        # validate the new_init_state if provided
 593        if new_init_state is not None:
 594            self._set_init_state(new_init_state)
 595        # restart game, skipping to init_state
 596        with open(self.init_state, "rb") as f:
 597            self._pyboy.load_state(f)
 598
 599        self.reset_count += 1
 600        self.step_count = 0
 601        self.state_tracker.reset()
 602        self.video_writer.close_video()
 603        return
 604
 605    def get_current_frame(self) -> np.ndarray:
 606        """
 607        Renders the currently rendered screen of the emulator and returns it as a numpy array.
 608
 609        Returns:
 610            np.ndarray: The rendered image as a numpy array.
 611        """
 612        return self.state_parser.get_current_frame()
 613
 614    def _update_listeners_after_actions(self, frames: np.ndarray):
 615        """
 616        Updates the state tracker after a batch of actions are run on the emulator without step()
 617
 618        You should *not* call this method when implementing HighLevelActions, instead call step(), track the states at each step, and return the list of transition states.
 619
 620        Args:
 621            frames (np.ndarray): Frames of shape [n_frames, H, W, C] that contain the frames which elapsed during the run of the actions outside step
 622        """
 623        self.state_tracker.step(frames)
 624
 625    def _get_unique_frames(self, frames: np.ndarray) -> np.ndarray:
 626        """
 627        Removes duplicate frames from a stack of frames.
 628
 629        Args:
 630            frames (np.ndarray): Stack of frames of shape [n_frames, H, W, C].
 631        Returns:
 632            np.ndarray: Stack of frames with duplicates removed. Shape is [n_unique_frames, H, W, C].
 633        """
 634        unique_frames = []
 635        for frame in frames:
 636            if len(unique_frames) == 0:
 637                unique_frames.append(frame)
 638            else:
 639                if not np.array_equal(frame, unique_frames[-1]):
 640                    unique_frames.append(frame)
 641        unique_frames = np.stack(unique_frames, axis=0)
 642        return unique_frames
 643
 644    def step(self, action: LowLevelActions = None) -> Tuple[Optional[np.ndarray], bool]:
 645        """
 646        Takes a step in the environment by performing the given action on the emulator. If saving video, starts the video recording on the first step.
 647
 648        :param action: Lowest level action to perform on the emulator.
 649        :type action: LowLevelActions
 650        :return:
 651            - The stack of frames that passed while performing the action, if rendering is enabled. Is of shape [n_frames (3 right now), height, width, channels]. Otherwise, None.
 652
 653            - Is max steps reached.
 654        :rtype: Tuple[Optional[np.ndarray], bool]
 655        """
 656        if action is not None:
 657            if action not in LowLevelActions:
 658                log_error(
 659                    f"Invalid action {action}. Must be one of {list(LowLevelActions)} or None",
 660                    self._parameters,
 661                )
 662        if self.step_count >= self.max_steps:
 663            log_warn(
 664                "Step called after max_steps reached. Please reset the environment.",
 665                self._parameters,
 666            )
 667            # This does not exit because some HighLevelActions may call step() multiple times in their execution.
 668            # It is not the best practice to allow this to happen, but it is easier to not error out here than check in every HighLevelAction, and this won't advantage the agent too much.
 669            # One consequence, however, is that max_steps then becomes a soft limit rather than a hard limit.
 670
 671        if self.save_video and self.step_count == 0:
 672            self.video_writer.start_video()
 673
 674        frames = self.run_action_on_emulator(action)
 675        self.step_count += 1
 676        frames = self._get_unique_frames(frames)
 677        self._update_listeners_after_actions(frames)
 678        return frames, self.check_if_done()
 679
 680    def get_state_parser(self) -> StateParser:
 681        """
 682        Returns the current game state parser instance.
 683
 684        Returns:
 685            StateParser: The current game state parser.
 686        """
 687        return self.state_parser
 688
 689    def run_action_on_emulator(
 690        self, action: LowLevelActions = None
 691    ) -> Optional[np.ndarray]:
 692        """
 693
 694        Performs the given action on the emulator by pressing and releasing the corresponding button.
 695
 696        Args:
 697            action (LowLevelActions): Lowest level action to perform on the emulator.
 698        Returns:
 699            Optional[np.ndarray]: The stack of frames that passed while performing the actions. Is of shape [n_frames (3 right now), height, width, channels]. Otherwise, None.
 700        """
 701        frames = None
 702        if action is not None:
 703            frames = []
 704            self._pyboy.send_input(action.value)
 705            press_step = self.press_step
 706            self._pyboy.tick(press_step, True)
 707            frames.append(self.get_current_frame())
 708            self._pyboy.send_input(ReleaseActions.release_actions.value[action])
 709            self._pyboy.tick(self.wait_ticks + press_step + 1, True)
 710            frames.append(self.get_current_frame())
 711            # self._pyboy.tick(1, True)
 712            # frames.append(self.get_current_frame())
 713            frames = np.stack(frames, axis=0)
 714        else:
 715            self._pyboy.tick(self.wait_ticks, True)
 716            frames = [self.get_current_frame()]
 717            frames = np.array(frames)
 718        if self.save_video and self.video_writer.video_running:
 719            self.video_writer.add_video_frames(frames, pressed_button=action)
 720        return frames
 721
 722    def check_if_done(self):
 723        """
 724        Checks if the max_steps limit has been reached.
 725        """
 726        done = self.step_count >= self.max_steps - 1
 727        return done
 728
 729    def close(self) -> StateTracker:
 730        """
 731        Closes the emulator and any associated resources.
 732        If the session directory is empty after closing, it will be deleted.
 733        """
 734        self.state_tracker.close()
 735        self._pyboy.stop(save=False)
 736        self.video_writer.close_video()
 737        self.state_tracker.close()
 738        # check if session directory is empty, and if so delete it
 739        if (
 740            os.path.exists(self.session_path)
 741            and len(os.listdir(self.session_path)) == 0
 742        ):
 743            os.rmdir(self.session_path)
 744        return self.state_tracker
 745
 746    def human_play(self, max_steps: int = None):
 747        """
 748        Allows a human to play the emulator using keyboard inputs.
 749        Args:
 750            max_steps (int, optional): Maximum number of steps to play. Defaults to gameboy_hard_max_steps in configs.
 751        """
 752        if max_steps is None:
 753            max_steps = self._parameters["gameboy_hard_max_steps"]
 754        log_info(
 755            "Starting human play mode. Use arrow keys and A(a)/B(s)/Start(enter) buttons to play. Close the window to exit.",
 756            self._parameters,
 757        )
 758        if self.headless:
 759            log_error(
 760                "Human play mode requires headless=False. Change the initialization",
 761                self._parameters,
 762            )
 763        self.reset()
 764        while True:
 765            self._pyboy.tick(1, True)
 766            self.state_tracker.step()
 767            if self.step_count >= max_steps:
 768                break
 769        self.close()
 770
 771    def random_play(self, max_steps: int = None):
 772        """
 773        Allows the emulator to play itself using (sort of) random actions.
 774        Args:
 775            max_steps (int, optional): Maximum number of steps to play. Defaults to gameboy_hard_max_steps in configs.
 776        """
 777        if max_steps is None:
 778            max_steps = self._parameters["gameboy_random_play_max_steps"]
 779        log_info("Starting random play mode.", self._parameters)
 780        self.reset()
 781        pbar = tqdm(total=max_steps, desc="Random Play Steps")
 782        allowed_actions = list(LowLevelActions)
 783        # remove the Start and Select actions from allowed actions to avoid menu spamming.
 784        allowed_actions.remove(LowLevelActions.PRESS_BUTTON_START)
 785        while self.step_count < max_steps:
 786            action = np.random.choice(allowed_actions)
 787            frames, done = self.step(action)
 788            pbar.update(1)
 789            if done:
 790                break
 791        pbar.close()
 792        self.close()
 793        log_info("Random play mode ended.", self._parameters)
 794
 795    def _dev_play(self, max_steps: int = None):
 796        """
 797        Allows a human to play the emulator using keyboard inputs. Does not route through step function.
 798        This function continuously reads from the parameters in the configs directory and if it detects a change in the `gameboy_dev_play_stop` parameter, will enter a breakpoint
 799
 800        Args:
 801            max_steps (int, optional): Maximum number of steps to play. Defaults to gameboy_hard_max_steps in configs.
 802        """
 803        if not hasattr(self.state_parser, "rom_data_path"):
 804            log_error(
 805                "Development play mode requires a StateParser with rom_data_path attribute.",
 806                self._parameters,
 807            )
 808        if max_steps is None:
 809            max_steps = self._parameters["gameboy_hard_max_steps"]
 810        log_info(
 811            "Starting human play mode. Use arrow keys and A(a)/B(s)/Start(enter) buttons to play. Close the window to exit. Open configs/gameboy_vars.yaml and set gameboy_dev_play_stop to true to enable development mode.",
 812            self._parameters,
 813        )
 814        if self.headless:
 815            log_error(
 816                "Human play mode requires headless=False. Change the initialization",
 817                self._parameters,
 818            )
 819        self.reset()
 820        valid_regions = []
 821        unassigned_regions = []
 822        for region_name, region in self.state_parser.named_screen_regions.items():
 823            valid_regions.append(region_name)
 824            if region.multi_targets is None:
 825                if region.target is None:
 826                    unassigned_regions.append(region_name)
 827            else:
 828                for target_name in region.multi_targets.keys():
 829                    if region.multi_targets[target_name] is None:
 830                        unassigned_regions.append((region_name, target_name))
 831        if len(unassigned_regions) > 0:
 832            log_warn(
 833                f"Unassigned regions (target array not set) are: {unassigned_regions}",
 834                self._parameters,
 835            )
 836        if self.save_video:
 837            self.video_writer.start_video()
 838        while True:
 839            self._parameters = load_parameters()
 840            if not self._parameters["gameboy_dev_play_stop"]:
 841                self._pyboy.tick(1, True)
 842                self.state_tracker.step()
 843                frames = [self.get_current_frame()]
 844                frames = np.array(frames)
 845                if self.save_video:
 846                    self.video_writer.add_video_frames(frames, show_button=False)
 847            else:
 848                tracker_report = self.state_tracker.report()
 849                tracker_report["core"].pop("current_frame", None)
 850                tracker_report["core"].pop("passed_frames", None)
 851                dev_instructions = f"""
 852                In development mode.
 853                Enter 'e' to close the emulator.
 854                Enter '' to re-enter normal play mode (remember to change gameboy_dev_play_stop back to false in configs or it'll stop again). 
 855                Enter 'p' to print the current state.
 856                Enter 'w' to pass a single tick without any action.
 857                Enter 's <state_name>' to save the current state as a .state file.
 858                Enter 'l <state_name>' to load a .state file.
 859                Enter 'c <region_name> <save_name / None if region.target_path is set>' to capture a named region and save it as a .npy file. To enter a multi-target region use the format "c <region_name>,<target_name> <save_name>" (no spaces in between region and target name)
 860                Enter 'd <None / region_name>' to draw a named region and display the current screen with the region drawn.
 861                Enter 'b' to enter a breakpoint.
 862                Enter 'g <code>' to apply a GameShark code (e.g. g 0101C7C9).
 863                Valid region names are: {valid_regions}
 864                Initially unassigned regions (target array not set) were: {unassigned_regions}\n\t Note: This list does not update as you assign targets during this session.
 865                Current State: 
 866                """
 867                log_info(dev_instructions, self._parameters)
 868                tracker_report.pop("ocr", None)
 869                log_dict(tracker_report, parameters=self._parameters)
 870                user_input = input("Dev mode input: ")
 871                user_input = user_input.lower().strip()
 872                first_char = user_input[0] if len(user_input) > 0 else ""
 873                allowed_inputs = ["e", "", "p", "w", "c", "s", "l", "d", "b", "g"]
 874                if first_char not in allowed_inputs:
 875                    log_warn(
 876                        f"Invalid input {user_input}. Valid inputs are: {allowed_inputs}",
 877                        self._parameters,
 878                    )
 879                    continue
 880                if first_char == "e":
 881                    log_info("Exiting human play mode.", self._parameters)
 882                    break
 883                elif first_char == "":
 884                    log_info(
 885                        "Exiting development mode. Resuming normal play.",
 886                        self._parameters,
 887                    )
 888                    continue
 889                elif first_char == "p":
 890                    log_info(
 891                        f"Current State:\n{str(self.state_tracker)}", self._parameters
 892                    )
 893                    continue
 894                elif first_char == "w":
 895                    self._pyboy.tick(1, True)
 896                    self.state_tracker.step()
 897                    continue
 898                elif first_char == "s" or first_char == "l":
 899                    parts = user_input.split(" ")
 900                    if len(parts) != 2:
 901                        log_warn(f"Invalid input {user_input}.", self._parameters)
 902                        continue
 903                    state_name = parts[1]
 904                    if not state_name.endswith(".state"):
 905                        state_name = state_name + ".state"
 906                    state_path = os.path.join(
 907                        self.state_parser.rom_data_path, "states", state_name
 908                    )
 909                    if first_char == "s":
 910                        if os.path.exists(state_path):
 911                            confirm_input = input(
 912                                f"State file {state_path} already exists. Overwrite? (y/n): "
 913                            )
 914                            if confirm_input.lower().strip() != "y":
 915                                log_info("Aborting save state.", self._parameters)
 916                                continue
 917                        self.save_state(state_path)
 918                    else:
 919                        if not os.path.exists(state_path):
 920                            log_warn(
 921                                f"State file {state_path} does not exist. Cannot load.",
 922                                self._parameters,
 923                            )
 924                            continue
 925                        self.set_init_state(state_path)
 926                elif first_char == "b":
 927                    grid_cells = self.state_parser.capture_grid_cells(
 928                        current_frame=self.get_current_frame(), y_offset=0
 929                    )
 930                    breakpoint()
 931                elif first_char == "g":
 932                    parts = user_input.split(" ")
 933                    if len(parts) != 2:
 934                        log_warn(
 935                            f"Invalid input {user_input}. Usage: g <code>",
 936                            self._parameters,
 937                        )
 938                    else:
 939                        code = parts[1].upper()
 940                        self._pyboy.gameshark.add(code)
 941                        log_info(f"GameShark code applied: {code}", self._parameters)
 942                    continue
 943                else:
 944                    current_frame = self.get_current_frame()
 945                    # draw it even if c, so we can see what we're capturing
 946                    save_path = None
 947                    if first_char == "c":
 948                        parts = user_input.split(" ")
 949                        if len(parts) != 3:
 950                            if len(parts) != 2:
 951                                log_warn(
 952                                    f"Invalid input {user_input}.", self._parameters
 953                                )
 954                                continue
 955                            else:
 956                                region_name = parts[1].split(",")[0]
 957                                region = self.state_parser.named_screen_regions[
 958                                    region_name
 959                                ]
 960                                if region.multi_targets is None:
 961                                    save_path = region.target_path
 962                                else:
 963                                    if "," not in parts[1]:
 964                                        log_warn(
 965                                            f"Region {region_name} is a multi-target region. Specify target",
 966                                            self._parameters,
 967                                        )
 968                                        continue
 969                                    target_name = parts[1].split(",")[1]
 970                                    if target_name not in region.multi_target_paths:
 971                                        log_warn(
 972                                            f"Target name {target_name} not found in region {region_name} with targets {region.multi_targets.keys()}.",
 973                                            self._parameters,
 974                                        )
 975                                        continue
 976                                    save_path = region.multi_target_paths[target_name]
 977                                if save_path is None:
 978                                    log_warn(
 979                                        f"Region {region_name} does not have a target path specified. Please provide a save name.",
 980                                        self._parameters,
 981                                    )
 982                                    continue
 983                        else:
 984                            save_name = parts[2]
 985                            if not save_name.endswith(".npy"):
 986                                save_name = save_name + ".npy"
 987                                save_path = os.path.join(
 988                                    self.state_parser.rom_data_path,
 989                                    "captures",
 990                                    save_name,
 991                                )
 992                        file_makedir(save_path)
 993                    elif first_char == "d":
 994                        parts = user_input.split(" ")
 995                        if len(parts) == 2:
 996                            region_name = parts[1].split(",")[
 997                                0
 998                            ]  # Shouldn't need but anyway.
 999                        elif len(parts) == 1:
1000                            region_name = "Full Screen"
1001                        else:
1002                            log_warn(f"Invalid input {user_input}.", self._parameters)
1003                            continue
1004                    if region_name == "Full Screen":
1005                        drawn_frame = current_frame  # self.state_parser.draw_grid_overlay(current_frame)
1006                    else:
1007                        drawn_frame = self.state_parser.draw_named_region(
1008                            current_frame, region_name
1009                        )
1010                    plt.imshow(drawn_frame[:, :, 0], cmap="gray")
1011                    plt.title(f"Region: {region_name}")
1012                    plt.show()
1013                    if first_char == "c":
1014                        captured_region = self.state_parser.capture_named_region(
1015                            current_frame, region_name
1016                        )
1017                        plt.imshow(captured_region[:, :, 0], cmap="gray")
1018                        plt.title(f"Captured Region: {region_name}")
1019                        plt.show()
1020                        existing_file = os.path.exists(save_path)
1021                        existing_str = (
1022                            ""
1023                            if not existing_file
1024                            else " (will overwrite existing file)"
1025                        )
1026                        confirmation_input = input(
1027                            f"Save captured region {region_name} to {save_path}? (y/n) {existing_str}: "
1028                        )
1029                        if confirmation_input.lower().strip() != "y":
1030                            log_info("Aborting capture region.", self._parameters)
1031                            continue
1032                        np.save(save_path, captured_region)
1033                        log_info(
1034                            f"Saved captured region {region_name} to {save_path}",
1035                            self._parameters,
1036                        )
1037            if self.step_count >= max_steps:
1038                break
1039        tracker = self.close()
1040        log_info("Human play mode ended.", self._parameters)
1041        # log_dict(tracker.report_final(), parameters=self._parameters)
1042
1043    def save_state(self, state_name: str, error_if_exists: bool = False):
1044        """
1045        Saves the current state of the emulator to a .state file.
1046        Args:
1047            state_name (str): Name of the state file to save (with or without .state extension).
1048            error_if_exists (bool): Whether to raise an error if the state file already exists.
1049        """
1050        if not state_name.endswith(".state"):
1051            state_name = state_name + ".state"
1052        state_dir = os.path.abspath(self.state_parser.rom_data_path + "/states/")
1053        potential_state_dir = os.path.abspath(os.path.dirname(state_name))
1054        if state_dir != potential_state_dir:
1055            if potential_state_dir != os.path.abspath(""):
1056                log_error(
1057                    f"Tried to save state file to {state_name}, which is outside of the states directory {state_dir}. This is not allowed.",
1058                    self._parameters,
1059                )
1060            state_name = os.path.abspath(os.path.join(state_dir, state_name))
1061        else:
1062            pass
1063        if os.path.exists(state_name):
1064            if error_if_exists:
1065                log_error(
1066                    f"State file {state_name} already exists. Will not overwrite...",
1067                    self._parameters,
1068                )
1069        file_makedir(state_name)
1070        with open(state_name, "wb") as f:
1071            self._pyboy.save_state(f)
1072        log_info(f"Saved state to {state_name}", self._parameters)
1073
1074    def delete_state(self, state_name: str, error_if_not_exists: bool = False):
1075        """
1076        Deletes a .state file from the states directory.
1077
1078        Args:
1079            state_name (str): Name of the state file to delete (with or without .state extension).
1080            error_if_not_exists (bool): Whether to raise an error if the state file does not exist.
1081        """
1082        if not state_name.endswith(".state"):
1083            state_name = state_name + ".state"
1084        state_path = os.path.join(self.state_parser.rom_data_path, "states", state_name)
1085        if not os.path.exists(state_path):
1086            if error_if_not_exists:
1087                log_error(
1088                    f"State file {state_path} does not exist. Cannot delete.",
1089                    self._parameters,
1090                )
1091            else:
1092                log_warn(
1093                    f"State file {state_path} does not exist. Cannot delete.",
1094                    self._parameters,
1095                )
1096            return
1097        # if somehow state_path isn't in the states directory, error out to avoid deleting random files
1098        if os.path.abspath(os.path.dirname(state_path)) != os.path.abspath(
1099            self.state_parser.rom_data_path + "/states/"
1100        ):
1101            log_error(
1102                f"Tried to delete state file at {state_path}, which is outside of the states directory. This is not allowed.",
1103                self._parameters,
1104            )
1105        os.remove(state_path)
1106        log_info(f"Deleted state file {state_path}", self._parameters)
1107
1108    def _sav_to_state(self, sav_file: Optional[str], state_name: str):
1109        """
1110        Loads a .sav file into the emulator and saves the corresponding .state file.
1111        Use this if you want to manually create .sav files and convert them to .state files for use as initial states.
1112        Requires `_open_to_first_state` to be implemented in the subclass to get past opening menus.
1113
1114        Args:
1115            save_file (str or None): Path to the .sav file to load. If None, looks for a .sav file in the same directory as the ROM with the same base name.
1116            state_name (str): Name of the state to save. The .state file will be saved in the states directory of the rom_data_path with this name.
1117        """
1118        log_info(
1119            "Trying to find .sav file and convert to .state file. This is a breaking operation, so the program will terminate after its completion.",
1120            self._parameters,
1121        )
1122        if sav_file is not None:
1123            expected_sav = sav_file
1124        else:
1125            if ".gbc" in self._gb_path:
1126                expected_sav = self._gb_path.replace(".gbc", ".sav")
1127            else:
1128                expected_sav = self._gb_path.replace(".gb", ".sav")
1129        if not os.path.exists(expected_sav):
1130            log_error(
1131                f"Expected .sav file at {expected_sav} to convert to .state file, but it does not exist.",
1132                self._parameters,
1133            )
1134        if state_name is None or state_name == "":
1135            log_error(
1136                "You must provide a state_name to save the .state file.",
1137                self._parameters,
1138            )
1139        # copy the .sav file to self._gb_path.gb.ram file
1140        save_destination = self._gb_path.replace(".gb", ".gb.ram")
1141        shutil.copyfile(expected_sav, save_destination)
1142        self.close()
1143        self._pyboy = PyBoy(
1144            self._gb_path,
1145            window="null",
1146        )
1147        self._pyboy.set_emulation_speed(0)
1148        self._open_to_first_state()
1149        self.save_state(state_name, error_if_exists=True)
1150        self._pyboy.stop(save=False)
1151        log_info(
1152            "State saved successfully. Exiting now to avoid issues ...",
1153            self._parameters,
1154        )
1155        # remove the .gb.ram file
1156        if os.path.exists(save_destination):
1157            os.remove(save_destination)
1158        sys.exit(0)
1159
1160    def _open_to_first_state(self):
1161        """
1162        Presses buttons on the emulator to get past the opening menus and into the game itself.
1163        You don't really need to implement this method, but if you do, you can use _sav_to_state to create state files from mGBA sav files.
1164        """
1165        raise NotImplementedError
1166
1167    def get_env_variant(self) -> str:
1168        """
1169        Returns a string identifier for the particular environment variant being used.
1170
1171        :return: string name identifier of the particular env e.g. PokemonRed
1172        """
1173        return self.game
1174
1175
1176def bytes_to_padded_hex_string(integer_value):
1177    """
1178    Converts a bytes object into a padded, '0x'-prefixed hexadecimal string.
1179    """
1180    # 1. Convert the bytes object back into an integer
1181    # Assumes big-endian order for the example '0x00a' -> 10
1182    # 2. Format the integer into a string with padding and the '0x' prefix
1183    # The 'x' specifier for hex, '#' adds '0x', '04' pads to 4 hex characters total
1184    # (not including the '0x' prefix for simple formatters like this, but managing width)
1185
1186    # A robust approach to match your exact output '0x00a':
1187    # You generally want enough width for your bytes. b'\n' is 1 byte, 2 hex chars.
1188
1189    return f"0x{integer_value:04x}"  # {0:04x} pads to 4 digits specifically