gameboy_worlds.interface.environment

   1from abc import abstractmethod, ABC
   2from typing import Optional, Type, Dict, Any, List, Tuple
   3
   4from gameboy_worlds.utils import (
   5    load_parameters,
   6    log_error,
   7    log_info,
   8    log_warn,
   9    get_lowest_level_subclass,
  10    verify_parameters,
  11    log_dict,
  12    import_pygame,
  13)
  14
  15
  16from gameboy_worlds.emulation import Emulator, StateTracker, TestTrackerMixin
  17from gameboy_worlds.emulation.registry import (
  18    get_state_tracker_class,
  19    get_train_init_states,
  20)
  21from gameboy_worlds.interface.controller import Controller
  22from gameboy_worlds.interface.action import HighLevelAction
  23
  24import numpy as np
  25import gymnasium as gym
  26import warnings
  27from copy import deepcopy
  28import uuid
  29
  30warnings.filterwarnings(
  31    "ignore",
  32    message="pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html",
  33)
  34# This is to ignore deprecation warnings from pygame about pkg_resources
  35
  36
  37class Environment(gym.Env, ABC):
  38    """Base class for environments interfacing with the emulator."""
  39
  40    REQUIRED_EMULATOR = Emulator
  41    """ The highest level emulator that the environment can interface with. """
  42
  43    REQUIRED_STATE_TRACKER = StateTracker
  44    """ The state tracker that tracks the minimal state information required for the environment to function. """
  45
  46    @staticmethod
  47    def override_emulator_kwargs(emulator_kwargs: dict) -> dict:
  48        """
  49        Override default emulator keyword arguments for this environment.
  50
  51        Override this method in subclasses to modify the default emulator keyword arguments.
  52
  53        You may want to use `override_state_tracker_class` or that style to ensure compatibility of state tracker classes.
  54
  55        Args:
  56            emulator_kwargs (dict): Incoming emulator keyword arguments.
  57        Returns:
  58            dict: The overridden emulator keyword arguments.
  59        """
  60        return emulator_kwargs
  61
  62    @staticmethod
  63    def override_state_tracker_class(
  64        emulator_kwargs: dict, required_state_tracker_class: Type[StateTracker]
  65    ):
  66        """
  67        Safely overrides the state tracker class for the environment.
  68
  69        Use this in `override_emulator_kwargs` to ensure that the lowest level state tracker class is chosen.
  70
  71        Args:
  72            emulator_kwargs (dict): Incoming emulator keyword arguments.
  73            required_state_tracker_class (Type[StateTracker]): Usually the required state tracker class for the environment.
  74        """
  75        game = emulator_kwargs["game"]
  76        has_option = "state_tracker_class" in emulator_kwargs
  77        incoming_state_tracker_class = emulator_kwargs.get(
  78            "state_tracker_class", "default"
  79        )
  80        if isinstance(incoming_state_tracker_class, str):
  81            incoming_state_tracker_class = get_state_tracker_class(
  82                game, incoming_state_tracker_class
  83            )
  84        if issubclass(incoming_state_tracker_class, required_state_tracker_class):
  85            return incoming_state_tracker_class
  86        elif issubclass(required_state_tracker_class, incoming_state_tracker_class):
  87            emulator_kwargs["state_tracker_class"] = required_state_tracker_class
  88        else:
  89            emulator_kwargs["state_tracker_class"] = (
  90                incoming_state_tracker_class  # Don't know which one to pick, so just go with the incoming one.
  91            )
  92        return
  93
  94    def __init__(
  95        self,
  96        emulator: Emulator,
  97        controller: Controller,
  98        parameters: Optional[dict] = None,
  99    ):
 100        """
 101        Ensures that the environment has the required attributes.
 102        All subclasses must call this __init__ method AFTER setting up the required attributes.
 103
 104        If you are implementing a subclass, ensure that the following attributes are set:
 105            - observation_space: gym space defining observation space structure
 106
 107        """
 108        self._parameters = load_parameters(parameters)
 109        self._emulator = emulator
 110        self._controller = controller
 111        required_attributes = ["observation_space"]
 112        for attr in required_attributes:
 113            if not hasattr(self, attr):
 114                log_error(
 115                    f"Environment requires attribute '{attr}' to be set. Implement this in the subclass __init__",
 116                    self._parameters,
 117                )
 118        self.observation_space: gym.spaces.Space = self.observation_space
 119        if not issubclass(type(self._emulator), self.REQUIRED_EMULATOR):
 120            log_error(
 121                f"Environment requires an Emulator of type {self.REQUIRED_EMULATOR}, but got {type(self._emulator)}",
 122                self._parameters,
 123            )
 124        if not isinstance(self._controller, Controller):
 125            log_error(
 126                f"Environment requires a Controller instance, but got {type(self._controller)}",
 127                self._parameters,
 128            )
 129        self.REQUIRED_STATE_TRACKER = get_lowest_level_subclass(
 130            [self.REQUIRED_STATE_TRACKER, self._controller.REQUIRED_STATE_TRACKER]
 131        )
 132        if not issubclass(
 133            type(self._emulator.state_tracker), self.REQUIRED_STATE_TRACKER
 134        ):
 135            log_error(
 136                f"Environment requires a StateTracker of type {self.REQUIRED_STATE_TRACKER}, but got {type(self._emulator.state_tracker)}",
 137                self._parameters,
 138            )
 139        self._controller.assign_emulator(self._emulator)
 140        self._rng = np.random.default_rng()
 141        self.action_space = self._controller.get_action_space()
 142        """ The Gym action Space provided by the controller. """
 143        self.actions = self._controller.ACTIONS
 144        """ A list of HighLevelAction Types provided by the controller. """
 145        self.render_mode = "human"
 146        """ The render mode of the environment. Supports 'human' and 'rgb_array', but strongly assumes 'human' as can just read the emulator screen from `get_info` """
 147        self._window = None
 148        """ The pygame window for rendering in 'human' mode. Initialized on first render call. """
 149        self._clock = None
 150        """ The pygame clock for rendering in 'human' mode. Initialized on first render call. """
 151        self.reset()  # I don't think this will cause issues, but should check that resetting here works well with gymnasium SyncVectorEnv final_obs construction.
 152
 153    def save_custom_state(self, state_name: str):
 154        """
 155        Saves a custom state of the emulator. This is useful for saving states during training or evaluation that can be loaded later for analysis or replay.
 156
 157        Args:
 158            state_name (str): Name of the state to save. This will be saved as a .state file in the states directory.
 159        """
 160        # don't allow path like state names
 161        if (
 162            "/" in state_name
 163            or "\\" in state_name
 164            or " " in state_name
 165            or not state_name.isalnum()
 166        ):
 167            log_error(
 168                f"State name '{state_name}' is invalid. State names must be alphanumeric and cannot contain spaces or path characters.",
 169                self._parameters,
 170            )
 171        state_name = state_name.replace(
 172            "custom_", ""
 173        )  # prevent users from accidentally adding the prefix and causing confusion about the actual saved state name.
 174        state_name = f"custom_{state_name}"
 175        self._emulator.save_state(state_name=state_name)
 176        return state_name
 177
 178    def delete_custom_state(self, state_name: str):
 179        """
 180        Deletes a custom state of the emulator that was previously saved with `save_custom_state`.
 181
 182        Args:
 183            state_name (str): Name of the state to delete. This should be the name returned by `save_custom_state`.
 184        """
 185        state_name = state_name.replace(
 186            "custom_", ""
 187        )  # prevent users from accidentally adding the prefix and causing confusion about the actual saved state name.
 188        state_name = f"custom_{state_name}"
 189        self._emulator.delete_state(state_name=state_name)
 190
 191    def load_custom_state(self, state_name: str):
 192        """
 193        Loads a custom state of the emulator that was previously saved with `save_custom_state`.
 194
 195        Args:
 196            state_name (str): Name of the state to load. This should be the name returned by `save_custom_state`.
 197        """
 198        state_name = f"custom_{state_name}"
 199        self._emulator.set_init_state(state_name)
 200        self.reset()  # reset to apply the new init state
 201
 202    @abstractmethod
 203    def get_observation(
 204        self,
 205        *,
 206        action: Optional[HighLevelAction] = None,
 207        action_kwargs: Optional[dict] = None,
 208        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 209        action_success: Optional[int] = None,
 210    ) -> gym.spaces.Space:
 211        """
 212        Returns the current observation from the emulator. Must match self.observation_space.
 213        Args:
 214            action (Optional[HighLevelAction]): The previous action taken.
 215            action_kwargs (dict): The keyword arguments used for the action.
 216            transition_states (Optional[List[Dict[str, Dict[str, Any]]]]): The states observed during the action execution.
 217            action_success (Optional[int]): The success code of the action.
 218
 219        Returns:
 220            observation (gym.spaces.Space): The current observation.
 221        """
 222        raise NotImplementedError
 223
 224    def get_info(
 225        self,
 226        *,
 227        action: Optional[HighLevelAction] = None,
 228        action_kwargs: Optional[dict] = None,
 229        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 230        action_success: Optional[int] = None,
 231    ) -> Dict[str, Dict[str, Any]]:
 232        """
 233        Returns the full state information as defined by the emulator's state tracker.
 234
 235        Creates additional fields:
 236        - "core"/"previous_action_details": A tuple of (action, action_kwargs, transition_states, action_success, action_return)
 237        - "core"/"transition_passed_frames": An array of all frames passed during the action execution
 238        - "ocr"/"transition_ocr_regions": A list of OCR regions captured during the action execution
 239
 240        :param action: HighLevelAction taken
 241        :type action: Optional[HighLevelAction]
 242        :param action_kwargs: Keyword arguments for the action
 243        :type action_kwargs: Optional[dict]
 244        :param transition_states: List of states observed during the action execution
 245        :type transition_states: Optional[List[Dict[str, Dict[str, Any]]]]
 246        :param action_success:  Success code of the action
 247        :type action_success: Optional[int]
 248        :return: Full state information from the state tracker
 249        :rtype: Dict[str, Dict[str, Any]]
 250        """
 251        state_info = self._emulator.state_tracker.report()
 252        if action is not None:  # then transition_states should not be empty
 253            # Attach the action details to the info
 254            last_state = transition_states[-1]
 255            if "action_return" in last_state["core"]:
 256                action_return = last_state["core"]["action_return"]
 257            else:
 258                action_return = None
 259            state_info["core"]["previous_action_details"] = (
 260                action,
 261                action_kwargs,
 262                transition_states,
 263                action_success,
 264                action_return,
 265            )
 266
 267            # Aggregate passed frames from transition states
 268            all_passed_frames = transition_states[0]["core"]["passed_frames"]
 269            for transition_state in transition_states[1:]:
 270                all_passed_frames = np.concatenate(
 271                    [all_passed_frames, transition_state["core"]["passed_frames"]],
 272                    axis=0,
 273                )
 274            state_info["core"][
 275                "transition_passed_frames"
 276            ] = all_passed_frames  # Will include the current state info last frame as as the final entry
 277
 278            # Aggregate OCR texts from transition states
 279            all_ocr_regions = []
 280            for transition_state in transition_states:
 281                if (
 282                    "ocr" in transition_state
 283                    and "ocr_regions" in transition_state["ocr"]
 284                ):
 285                    all_ocr_regions.append(transition_state["ocr"]["ocr_regions"])
 286            if "ocr" in state_info:
 287                state_info["ocr"]["transition_ocr_regions"] = all_ocr_regions
 288            else:
 289                state_info["ocr"] = {"transition_ocr_regions": all_ocr_regions}
 290        return state_info
 291
 292    def get_final_info(self) -> Dict[str, Dict[str, Any]]:
 293        """
 294        Returns the final state information from the emulator when all episodes are done.
 295        Will involve summaries over all episodes played.
 296        Returns:
 297            info (dict): The final state information from the state tracker.
 298        """
 299        return self._emulator.state_tracker.report_final()
 300
 301    def reset(
 302        self, *, seed: Optional[int] = None, options: Optional[dict] = None
 303    ) -> Tuple[gym.spaces.Space, Dict[str, Dict[str, Any]]]:
 304        """
 305        Resets the environment and emulator to the initial state.
 306        Args:
 307            seed (int, optional): Seed for random number generators.
 308            options (dict, optional): Additional options for resetting the environment.
 309        Returns:
 310            observation (object): The initial observation of the environment.
 311
 312            info (dict): Additional information about the reset.
 313        """
 314        super().reset(seed=seed, options=options)
 315        self._emulator.reset()
 316        self.seed(seed)
 317        observation, info = self.get_observation(), self.get_info()
 318        return observation, info
 319
 320    @abstractmethod
 321    def determine_reward(
 322        self,
 323        start_state: Dict[str, Dict[str, Any]],
 324        *,
 325        action: Optional[HighLevelAction] = None,
 326        action_kwargs: Optional[dict] = None,
 327        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 328        action_success: Optional[int] = None,
 329    ) -> float:
 330        """
 331        Determines the reward based on the transition from start_state through transition_states.
 332        Args:
 333            start_state (Dict[str, Dict[str, Any]]): The state before the action was taken.
 334            action (HighLevelAction): The HighLevelAction action taken.
 335            action_kwargs (dict): The keyword arguments used for the action.
 336            transition_states (List[Dict[str, Dict[str, Any]]]): A list of states observed during the action execution.
 337            action_success (bool): Whether the action was successful.
 338        Returns:
 339            float: The computed reward.
 340        """
 341        raise NotImplementedError
 342
 343    def determine_truncated(
 344        self,
 345        start_state: Dict[str, Dict[str, Any]],
 346        *,
 347        action: Optional[HighLevelAction] = None,
 348        action_kwargs: Optional[dict] = None,
 349        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 350        action_success: Optional[int] = None,
 351    ) -> bool:
 352        """
 353        Determines whether the episode playthrough has exceeded some maximum step count or other truncation criteria based on the transition from start_state through transition_states.
 354        This method is can be overidden to implement custom truncation logic, but it must always return:
 355        `super().determine_truncated() or <custom_truncation_logic_bool>`
 356
 357        Args:
 358            start_state (Dict[str, Dict[str, Any]]): The state before the action was taken.
 359            action (HighLevelAction): The HighLevelAction action taken.
 360            action_kwargs (dict): The keyword arguments used for the action.
 361            transition_states (List[Dict[str, Dict[str, Any]]]): A list of states observed during the action execution.
 362            action_success (bool): Whether the action was successful.
 363        Returns:
 364            bool: Whether the episode is terminated.
 365        """
 366        return self._emulator.check_if_done()
 367
 368    @abstractmethod
 369    def determine_terminated(
 370        self,
 371        start_state: Dict[str, Dict[str, Any]],
 372        *,
 373        action: Optional[HighLevelAction] = None,
 374        action_kwargs: Optional[dict] = None,
 375        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 376        action_success: Optional[int] = None,
 377    ) -> bool:
 378        """
 379        Determines whether the episode reaches the goal / terminal state based on the transition from start_state through transition_states.
 380        This method is NOT meant to be used to determine if the step count has exceeded the maximum.
 381
 382        Args:
 383            start_state (Dict[str, Dict[str, Any]]): The state before the action was taken.
 384            action (HighLevelAction): The HighLevelAction action taken.
 385            action_kwargs (dict): The keyword arguments used for the action.
 386            transition_states (List[Dict[str, Dict[str, Any]]]): A list of states observed during the action execution.
 387            action_success (bool): Whether the action was successful.
 388        Returns:
 389            bool: Whether the episode is terminated.
 390        """
 391        pass
 392
 393    def before_step(self, action: Type[HighLevelAction], action_kwargs: dict):
 394        """
 395        Implement any logic that needs to be executed before each step in the environment.
 396        """
 397        return
 398
 399    def after_step(
 400        self,
 401        start_state: Dict[str, Dict[str, Any]],
 402        action: Type[HighLevelAction],
 403        action_kwargs: dict,
 404        transition_states: List[Dict[str, Dict[str, Any]]],
 405        action_success: int,
 406    ):
 407        """
 408        Implement any logic that needs to be executed after each step in the environment.
 409
 410        Args:
 411            start_state (Dict[str, Dict[str, Any]]): The state before the action was taken.
 412            action (HighLevelAction): The HighLevelAction action taken.
 413            action_kwargs (dict): The keyword arguments used for the action.
 414            transition_states (List[Dict[str, Dict[str, Any]]]): A list of states observed during the action execution.
 415            action_success (int): Whether the action was successful.
 416        """
 417        return
 418
 419    def step(
 420        self, action: gym.spaces.OneOf
 421    ) -> Tuple[gym.spaces.Space, float, bool, bool, Dict[str, Dict[str, Any]]]:
 422        """
 423        Executes the given Gym Space action in the environment via the controller.
 424        Use step_high_level_action to execute high level actions directly.
 425
 426        Args:
 427            action (gym.spaces.OneOf): The action to execute. Must be a valid action in the controller's action space.
 428
 429        Returns:
 430            observation (gym.spaces.Space): The observation after executing the action.
 431            reward (float): The reward obtained from executing the action.
 432            terminated (bool): Whether the episode has ended (reached the terminal state of the MDP).
 433            truncated (bool): Whether the episode was truncated (exceeded the maximum allowed steps).
 434            info (Dict[str, Dict[str, Any]]): Full state information.
 435        """
 436        high_level_action, kwargs = self._controller._space_action_to_high_level_action(
 437            action
 438        )
 439        return self.step_high_level_action(high_level_action, **kwargs)
 440
 441    def step_high_level_action(
 442        self, action: Type[HighLevelAction], **kwargs
 443    ) -> Tuple[gym.spaces.Space, float, bool, bool, Dict[str, Dict[str, Any]]]:
 444        """
 445        Executes the given High Level action in the environment via the controller.
 446        If the action is invalid according to the controller, will not perform any action and will simply return the current observation, a reward of 0, and terminated and truncated as False. The info will also include a field "invalid_action"=True to indicate that the action was invalid.
 447
 448        :param action: The high level action class to execute.
 449        :type action: Type[HighLevelAction]
 450        :param kwargs: Additional arguments required for the specific high level action.
 451        :type kwargs: Dict[str, Any]
 452        :return:
 453            - observation (gym.spaces.Space): The observation after executing the action.
 454
 455            - reward (float): The reward obtained from executing the action.
 456
 457            - terminated (bool): Whether the episode has ended (reached the terminal state of the MDP).
 458
 459            - truncated (bool): Whether the episode was truncated (exceeded the maximum allowed steps).
 460
 461            - info (Dict[str, Dict[str, Any]]): Full state information.
 462        :rtype: Tuple[Space, float, bool, bool, Dict[str, Dict[str, Any]]]
 463        """
 464        if self._emulator.check_if_done():
 465            log_error(
 466                "Cannot step environment because emulator indicates done. Please reset the environment.",
 467                self._parameters,
 468            )
 469        start_state = self.get_info()
 470        self.before_step(action, kwargs)
 471        transition_states, action_success = self._controller.execute(action, **kwargs)
 472        if (
 473            transition_states is None
 474        ):  # then the action was not a valid one according to the controller.
 475            observation = self.get_observation()
 476            current_state = self.get_info()
 477            terminated = self.determine_terminated(start_state=start_state)
 478            truncated = self.determine_truncated(start_state=start_state)
 479            reward = self.determine_reward(start_state=start_state) - abs(
 480                self._parameters["invalid_action_penalty"]
 481            )
 482            current_state["invalid_action"] = True
 483            return observation, reward, terminated, truncated, current_state
 484        self.after_step(start_state, action, kwargs, transition_states, action_success)
 485        truncated = self.determine_truncated(
 486            start_state=start_state,
 487            action=action,
 488            action_kwargs=kwargs,
 489            transition_states=transition_states,
 490            action_success=action_success,
 491        )
 492
 493        observation = self.get_observation(
 494            action=action,
 495            action_kwargs=kwargs,
 496            transition_states=transition_states,
 497            action_success=action_success,
 498        )
 499        current_state = self.get_info(
 500            action=action,
 501            action_kwargs=kwargs,
 502            transition_states=transition_states,
 503            action_success=action_success,
 504        )
 505        terminated = self.determine_terminated(
 506            start_state=start_state,
 507            action=action,
 508            action_kwargs=kwargs,
 509            transition_states=transition_states,
 510            action_success=action_success,
 511        )
 512
 513        reward = self.determine_reward(
 514            start_state=start_state,
 515            action=action,
 516            action_kwargs=kwargs,
 517            transition_states=transition_states,
 518            action_success=action_success,
 519        )
 520        return observation, reward, terminated, truncated, current_state
 521
 522    def step_str(
 523        self, input_str: str
 524    ) -> Tuple[gym.spaces.Space, float, bool, bool, Dict[str, Dict[str, Any]]]:
 525        """
 526        Attempts to execute an input string representation of an action.
 527        Useful for human play or VLM interaction.
 528        If the action is an invalid string, will not perform any action and will simply return Nones.
 529
 530        :param input_str: The input string representing the action.
 531        :type input_str: str
 532        :return:
 533            - observation (gym.spaces.Space): The observation after executing the action.
 534
 535            - reward (float): The reward obtained from executing the action.
 536
 537            - terminated (bool): Whether the episode has ended (reached the terminal state of the MDP).
 538
 539            - truncated (bool): Whether the episode was truncated (exceeded the maximum allowed steps).
 540
 541            - info (Dict[str, Dict[str, Any]]): Full state information.
 542        :rtype: Tuple[Space, float, bool, bool, Dict[str, Dict[str, Any]]]
 543        """
 544        action, kwargs = self.string_to_high_level_action(input_str)
 545        if (
 546            action is None
 547        ):  # not a valid action, will not perform an action and will simply return Nones.
 548            return None, None, None, None, None
 549        return self.step_high_level_action(action, **kwargs)
 550
 551    def string_to_high_level_action(
 552        self, input_str: str
 553    ) -> Tuple[Optional[Type[HighLevelAction]], Optional[dict]]:
 554        """
 555        Attempts to convert an input string representation of an action into a HighLevelAction and its parameters.
 556        Useful for human play or VLM interaction.
 557
 558        :param input_str: The input string representing the action.
 559        :type input_str: str
 560        :return: A tuple containing the HighLevelAction class and its execution parameters dictionary. If the input string is invalid, returns (None, None).
 561        :rtype: Tuple[Type[HighLevelAction] | None, dict | None]
 562        """
 563        return self._controller.string_to_high_level_action(input_str)
 564
 565    def _simulate(self, step_fn, *args, **kwargs):
 566        """
 567        Executes step_fn(*args, **kwargs) without permanently advancing state.
 568
 569        Saves the emulator's current state, runs the step, then restores both the
 570        emulator's runtime state and its init_state pointer, and cleans up the
 571        temporary save file.
 572        """
 573        original_init_state = self._emulator.init_state
 574        tmp_name = uuid.uuid4().hex
 575        self.save_custom_state(tmp_name)
 576        try:
 577            result = step_fn(*args, **kwargs)
 578        finally:
 579            self.load_custom_state(tmp_name)
 580            self._emulator.init_state = original_init_state
 581            self.delete_custom_state(tmp_name)
 582        return result
 583
 584    def sim(
 585        self, action: gym.spaces.OneOf
 586    ) -> Tuple[gym.spaces.Space, float, bool, bool, Dict[str, Dict[str, Any]]]:
 587        """Like `step` but reverts the emulator to its pre-step state afterward.
 588
 589        .. warning::
 590            Because reversion requires a full emulator reset, all state tracker counters
 591            and accumulated metrics (e.g. steps taken, episode rewards) will be reset as
 592            a side effect. The returned info reflects the simulated step, not post-reset state.
 593        """
 594        return self._simulate(self.step, action)
 595
 596    def sim_str(
 597        self, input_str: str
 598    ) -> Tuple[gym.spaces.Space, float, bool, bool, Dict[str, Dict[str, Any]]]:
 599        """Like `step_str` but reverts the emulator to its pre-step state afterward.
 600
 601        .. warning::
 602            Because reversion requires a full emulator reset, all state tracker counters
 603            and accumulated metrics (e.g. steps taken, episode rewards) will be reset as
 604            a side effect. The returned info reflects the simulated step, not post-reset state.
 605        """
 606        return self._simulate(self.step_str, input_str)
 607
 608    def sim_high_level_action(
 609        self, action: Type[HighLevelAction], **kwargs
 610    ) -> Tuple[gym.spaces.Space, float, bool, bool, Dict[str, Dict[str, Any]]]:
 611        """Like `step_high_level_action` but reverts the emulator to its pre-step state afterward.
 612
 613        .. warning::
 614            Because reversion requires a full emulator reset, all state tracker counters
 615            and accumulated metrics (e.g. steps taken, episode rewards) will be reset as
 616            a side effect. The returned info reflects the simulated step, not post-reset state.
 617        """
 618        return self._simulate(self.step_high_level_action, action, **kwargs)
 619
 620    def close(self):
 621        """
 622        Closes the environment and the underlying emulator.
 623        """
 624        log_info("Closing environment and emulator.", self._parameters)
 625        self._emulator.close()
 626
 627    def _screen_render(self, screen: np.ndarray):
 628        """
 629        Renders the given screen using pygame in human mode.
 630        Args:
 631            screen (np.ndarray): The screen to render.
 632
 633        """
 634        pygame = import_pygame(self._parameters)
 635        if self._window is None:
 636            pygame.init()
 637            pygame.display.init()
 638            self._window = pygame.display.set_mode(
 639                (self._emulator.screen_shape[0], self._emulator.screen_shape[1])
 640            )
 641        if self._clock is None:
 642            self._clock = pygame.time.Clock()
 643        rgb = np.stack([screen[:, :, 0], screen[:, :, 0], screen[:, :, 0]], axis=2)
 644        pygame.surfarray.blit_array(self._window, rgb.swapaxes(0, 1))
 645        pygame.display.flip()
 646        self._clock.tick(60)  # Limit to 60 FPS
 647
 648    def render(self) -> Optional[np.ndarray]:
 649        """
 650        Gets the current screen from the emulator and renders it.
 651
 652        Use this method only if you want to generally run the emulator in headless mode but still want to see the screen occasionally.
 653
 654        Do not call this method if the emulator is not headless, you should already have a PyBoy interactive window open in that case.
 655
 656        Returns:
 657            If render_mode is 'rgb_array', returns the current screen as a numpy array. However this is always accessible via self.get_info()['core']['current_frame'], so this is mostly for Gym compatibility.
 658        """
 659        if self._emulator.headless == False:
 660            log_error(
 661                "You probably don't want to call render() when the emulator is not headless.",
 662                self._parameters,
 663            )
 664        screen = self._emulator.get_current_frame()  # shape: 144, 160, 1
 665        if self.render_mode == "human":
 666            self._screen_render(screen)
 667        elif self.render_mode == "rgb_array":
 668            return screen
 669        else:
 670            log_error(f"Unsupported render mode: {self.render_mode}", self._parameters)
 671
 672    def seed(self, seed: Optional[int] = None):
 673        """
 674        Seeds the environment's random number generator and the controller's RNG.
 675
 676        Args:
 677            seed (int, optional): The seed value.
 678        """
 679        self._controller.seed(seed)
 680        self._rng = np.random.default_rng(seed)
 681
 682    def render_obs(
 683        self,
 684        *,
 685        action: Optional[Type[HighLevelAction]] = None,
 686        action_kwargs: Optional[dict] = None,
 687        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 688        action_success: Optional[int] = None,
 689        action_return: Optional[Dict[str, Any]] = None,
 690    ):
 691        """
 692        Provide a way to render the output of `get_observation` to a human.
 693        Implement if you want to use the human_step_play method.
 694
 695        Args:
 696            action (Optional[Type[HighLevelAction]]): The previous action taken.
 697            action_kwargs (dict): The keyword arguments used for the action.
 698            transition_states (Optional[List[Dict[str, Dict[str, Any]]]]): The states observed during the action execution.
 699            action_success (Optional[int]): The success code of the action.
 700        """
 701        raise NotImplementedError
 702
 703    def render_info(
 704        self,
 705        *,
 706        action: Optional[Type[HighLevelAction]] = None,
 707        action_kwargs: Optional[dict] = None,
 708        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 709        action_success: Optional[int] = None,
 710        action_return: Optional[Dict[str, Any]] = None,
 711    ):
 712        """
 713        Provide a way to render the output of `get_info` to a human.
 714        Implement if you want to use the human_step_play method with show_info=True.
 715        Must always be a superset of render_obs (i.e. should also show what render_obs shows).
 716
 717        Args:
 718            action (Optional[Type[HighLevelAction]]): The previous action taken.
 719            action_kwargs (dict): The keyword arguments used for the action.
 720            transition_states (Optional[List[Dict[str, Dict[str, Any]]]]): The states observed during the action execution.
 721            action_success (Optional[int]): The success code of the action.
 722        """
 723        raise NotImplementedError
 724
 725    def get_action_strings(
 726        self, return_all: bool = False
 727    ) -> Dict[Type[HighLevelAction], str]:
 728        """
 729        Provide a way to verbalize the allowed high level actions, along with the format of the input parameters.
 730        Useful for prompting a VLM to choose an action.
 731
 732        :param return_all: If True, returns all possible actions and parameter formats. If False, returns only the actions that are valid in the current state.
 733        :type return_all: bool
 734        :return: A dictionary mapping high level actions to their verbalizations and input formats.
 735        :rtype: Dict[Type[HighLevelAction], str]
 736        """
 737        return self._controller.get_action_strings(return_all=return_all)
 738
 739    def human_step_play(
 740        self, max_steps: int = 50, show_obs: bool = True, show_info: bool = True
 741    ) -> Tuple[List[float], bool, bool]:
 742        """
 743        Opens a render window and allow the human to play through the environment as an agent would
 744
 745        Args:
 746            max_steps (int): max steps to take
 747            show_obs (bool): whether to show the observation space rendering after each step
 748            show_info (bool): whether to show the info rendering after each step. Will force show_obs=False if enabled.
 749
 750        Returns:
 751            rewards (List[float]): List of rewards obtained at each step.
 752            terminated (bool): Whether the episode has ended (reached the terminal state of the MDP).
 753            truncated (bool): Whether the episode was truncated (exceeded the maximum allowed emulator steps). It does not consider the max_steps parameter here.
 754        """
 755        observation, info = self.reset()
 756        self.render_mode = "human"
 757        log_info(f"Doing human step play for {max_steps} max steps...")
 758        steps = 0
 759        done = False
 760        terminated = False
 761        truncated = False
 762        rewards = []
 763        if show_info:
 764            self.render_info()
 765            show_obs = False
 766        if show_obs:
 767            self.render_obs()
 768        while not done and steps < max_steps:
 769            action_input = self._controller.get_action_strings()
 770            action_input_str = ""
 771            for action_class, action_str in action_input.items():
 772                action_input_str += f"- {action_str}\n"
 773            log_info(f"Allowed Actions: \n{action_input_str}", self._parameters)
 774            input_str = input("Enter Action: ").strip()
 775            (
 776                possible_obs,
 777                possible_reward,
 778                possible_terminated,
 779                possible_truncated,
 780                possible_info,
 781            ) = self.step_str(input_str)
 782            if possible_obs is not None:
 783                observation, reward, terminated, truncated, info = (
 784                    possible_obs,
 785                    possible_reward,
 786                    possible_terminated,
 787                    possible_truncated,
 788                    possible_info,
 789                )
 790                rewards.append(reward)
 791                if reward != 0:
 792                    log_info(f"Reward obtained: {reward}", self._parameters)
 793                (
 794                    action,
 795                    action_kwargs,
 796                    transition_states,
 797                    action_success,
 798                    action_return,
 799                ) = info["core"]["previous_action_details"]
 800                if show_info:
 801                    self.render_info(
 802                        action=action,
 803                        action_kwargs=action_kwargs,
 804                        transition_states=transition_states,
 805                        action_success=action_success,
 806                        action_return=action_return,
 807                    )
 808                if show_obs:
 809                    self.render_obs(
 810                        action=action,
 811                        action_kwargs=action_kwargs,
 812                        transition_states=transition_states,
 813                        action_success=action_success,
 814                        action_return=action_return,
 815                    )
 816            else:
 817                log_warn("That was not a valid input. did nothing", self._parameters)
 818            if terminated or truncated:
 819                log_info(
 820                    f"Episode finished! Terminated: {terminated}, Truncated: {truncated}",
 821                    self._parameters,
 822                )
 823                break
 824            steps += 1
 825            if steps >= max_steps:
 826                log_info(
 827                    f"Max steps {max_steps} reached. Ending episode.", self._parameters
 828                )
 829        return rewards, terminated, truncated
 830
 831
 832class DummyEnvironment(Environment):
 833    """A dummy environment that does nothing special."""
 834
 835    def __init__(
 836        self,
 837        emulator: Emulator,
 838        controller: Controller,
 839        parameters: Optional[dict] = None,
 840    ):
 841        """
 842        Initializes the DummyEnvironment with the given emulator and controller.
 843
 844        It is safe to overwrite the self.observation_space in the subclass after calling this __init__ method.
 845        """
 846        screen_shape = emulator.screen_shape
 847        self.observation_space = gym.spaces.Box(
 848            low=0, high=255, shape=(screen_shape[1], screen_shape[0]), dtype=np.uint8
 849        )
 850        """ The observation space is the raw pixel values of the emulator's screen. """
 851        super().__init__(
 852            emulator=emulator, controller=controller, parameters=parameters
 853        )
 854
 855    def get_observation(
 856        self,
 857        *,
 858        action=None,
 859        action_kwargs=None,
 860        transition_states=None,
 861        action_success=None,
 862    ):
 863        if transition_states is None:
 864            current_state = self.get_info()
 865            screen = current_state["core"]["current_frame"]
 866        else:
 867            screen = transition_states[-1]["core"]["current_frame"]
 868        return screen
 869
 870    def determine_reward(self, **kwargs):
 871        return 0.0
 872
 873    def determine_terminated(self, **kwargs):
 874        return False
 875
 876    def render_obs(
 877        self,
 878        action=None,
 879        action_kwargs=None,
 880        transition_states=None,
 881        action_success=None,
 882        action_return=None,
 883    ):  # Might cause issues if you try to render() as well
 884        """
 885        Renders the screen.
 886        """
 887        screen = self.get_observation()
 888        self._screen_render(screen)
 889
 890    def render_info(
 891        self,
 892        action=None,
 893        action_kwargs=None,
 894        transition_states=None,
 895        action_success=None,
 896        action_return=None,
 897    ):
 898        info = deepcopy(self.get_info())
 899        if transition_states is not None and len(transition_states) > 0:
 900            screens = transition_states[0]["core"]["passed_frames"]
 901            for transition_state in transition_states[1:]:
 902                screens = np.concatenate(
 903                    [screens, transition_state["core"]["passed_frames"]], axis=0
 904                )
 905        else:
 906            if "passed_frames" in info["core"]:
 907                screens = info["core"]["passed_frames"]
 908            else:
 909                screens = None
 910        if screens is None:
 911            screens = [info["core"]["current_frame"]]
 912        for screen in screens:
 913            self._screen_render(screen)
 914        info["core"].pop("current_frame")
 915        info["core"].pop("passed_frames")
 916        if "ocr" in info:
 917            info.pop("ocr")
 918        if "transition_passed_frames" in info["core"]:
 919            info["core"].pop("transition_passed_frames")
 920        if "previous_action_details" in info["core"]:
 921            info["core"]["previous_action_details"] = (
 922                info["core"]["previous_action_details"][:2]
 923                + info["core"]["previous_action_details"][3:]
 924            )  # remove transition states to avoid huge logs
 925        log_info("State: ", self._parameters)
 926        log_dict(info, parameters=self._parameters)
 927
 928
 929class TestEnvironmentMixin:
 930    """
 931    Mixin class for testing environments.
 932    Ensures the State Tracker used is a TestTrackerMixin and checks these for termination / truncation.
 933    """
 934
 935    REQUIRED_STATE_TRACKER = TestTrackerMixin
 936
 937    def determine_truncated(
 938        self,
 939        start_state: Dict[str, Dict[str, Any]],
 940        *,
 941        action: Optional[HighLevelAction] = None,
 942        action_kwargs: Optional[dict] = None,
 943        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 944        action_success: Optional[int] = None,
 945    ) -> bool:
 946        if not isinstance(self._emulator.state_tracker, TestTrackerMixin):
 947            log_error(
 948                "TestEnvironmentMixin requires the emulator's state tracker to be a TestTrackerMixin.",
 949                self._parameters,
 950            )  # we don't need to repeat this for determine_terminated since its done here.
 951        if transition_states is None:
 952            return self._emulator.check_if_done()
 953        any_truncated = False
 954        for state in transition_states:
 955            if state["termination_truncation"]["truncated"]:
 956                any_truncated = True
 957                break
 958        return any_truncated or self._emulator.check_if_done()
 959
 960    def determine_terminated(
 961        self,
 962        start_state: Dict[str, Dict[str, Any]],
 963        *,
 964        action: Optional[HighLevelAction] = None,
 965        action_kwargs: Optional[dict] = None,
 966        transition_states: Optional[List[Dict[str, Dict[str, Any]]]] = None,
 967        action_success: Optional[int] = None,
 968    ) -> bool:
 969        if transition_states is None:
 970            return False
 971        any_terminated = False
 972        for state in transition_states:
 973            if state["termination_truncation"]["terminated"]:
 974                any_terminated = True
 975                break
 976        return any_terminated
 977
 978
 979class TrainEnvironmentMixin:
 980    """
 981    Mixin class for training environments.
 982    Records the allowed initial states for training and random shuffles between them when resetting.
 983    """
 984
 985    def reset(
 986        self, *, seed: Optional[int] = None, options: Optional[dict] = None
 987    ) -> Tuple[gym.spaces.Space, Dict[str, Dict[str, Any]]]:
 988        """
 989
 990        Args:
 991            seed (int, optional): Seed for random number generators.
 992            options (dict, optional): Additional options for resetting the environment.
 993        Returns:
 994            observation (object): The initial observation of the environment.
 995
 996            info (dict): Additional information about the reset.
 997        """
 998        game = self._emulator.game
 999        init_states = get_train_init_states(game, parameters=self._parameters)
1000        choice = self._rng.choice(len(init_states))
1001        init_state = init_states[choice]
1002        self._emulator.set_init_state(init_state)
1003        return super().reset(seed=seed, options=options)