gameboy_worlds.interface.pokemon.environments

  1from typing import Optional, Dict, Any, List, Tuple
  2
  3from gymnasium import spaces
  4
  5from gameboy_worlds.emulation.pokemon.base_metrics import CorePokemonMetrics
  6from gameboy_worlds.utils import load_parameters, log_dict, log_info
  7from gameboy_worlds.emulation.pokemon.emulators import PokemonEmulator
  8from gameboy_worlds.emulation.pokemon.trackers import (
  9    PokemonRedStarterTracker,
 10    PokemonOCRTracker,
 11)
 12from gameboy_worlds.interface.environment import (
 13    DummyEnvironment,
 14    Environment,
 15    TestEnvironmentMixin,
 16    TrainEnvironmentMixin,
 17)
 18from gameboy_worlds.interface.controller import Controller
 19
 20import gymnasium as gym
 21import numpy as np
 22
 23
 24class PokemonEnvironment(DummyEnvironment):
 25    """
 26    A basic Pokemon Environment.
 27    """
 28
 29    REQUIRED_EMULATOR = PokemonEmulator
 30    REQUIRED_STATE_TRACKER = CorePokemonMetrics
 31
 32
 33class PokemonOCREnvironment(PokemonEnvironment):
 34    """
 35    A Pokemon Environment that includes OCR region captures and agent state.
 36    """
 37
 38    REQUIRED_STATE_TRACKER = PokemonOCRTracker
 39    REQUIRED_EMULATOR = PokemonEmulator
 40
 41    @staticmethod
 42    def override_emulator_kwargs(emulator_kwargs: dict) -> dict:
 43        Environment.override_state_tracker_class(
 44            emulator_kwargs, PokemonOCREnvironment.REQUIRED_STATE_TRACKER
 45        )
 46        return emulator_kwargs
 47
 48
 49class PokemonTestEnvironment(TestEnvironmentMixin, PokemonOCREnvironment):
 50    pass
 51
 52
 53class PokemonTrainEnvironment(TrainEnvironmentMixin, PokemonOCREnvironment):
 54    pass
 55
 56
 57class PokemonRedStarterChoiceEnvironment(PokemonOCREnvironment):
 58    """
 59    An environment which starts at the starter selection point in PokemonRed and terminates when the player selects a starter Pokemon.
 60    """
 61
 62    REQUIRED_TRACKER = PokemonRedStarterTracker
 63
 64    def override_emulator_kwargs(emulator_kwargs: dict) -> dict:
 65        """
 66        Override default emulator keyword arguments for this environment.
 67        """
 68        Environment.override_state_tracker_class(
 69            emulator_kwargs, PokemonRedStarterTracker
 70        )
 71        emulator_kwargs["init_state"] = "starter"
 72        return emulator_kwargs
 73
 74    def determine_terminated(
 75        self,
 76        start_state,
 77        *,
 78        action=None,
 79        action_kwargs=None,
 80        transition_states=None,
 81        action_success=None,
 82    ) -> bool:
 83        super_terminated = super().determine_terminated(
 84            start_state=start_state,
 85            action=action,
 86            action_kwargs=action_kwargs,
 87            transition_states=transition_states,
 88            action_success=action_success,
 89        )
 90        if transition_states is None:
 91            return super_terminated
 92        states = transition_states
 93        for state in states:
 94            starter_chosen = state["pokemon_red_starter"]["current_starter"]
 95            if starter_chosen is not None:
 96                return True
 97        return super_terminated
 98
 99
100class PokemonRedChooseCharmanderEnvironment(PokemonRedStarterChoiceEnvironment):
101    """
102    Reward the agent for choosing Charmander as quickly as possible.
103    """
104
105    STEP_WISE_REWARD = False
106
107    def determine_reward(
108        self,
109        start_state,
110        *,
111        action=None,
112        action_kwargs=None,
113        transition_states=None,
114        action_success=None,
115    ) -> float:
116        """
117        Reward the agent for choosing Charmander as quickly as possible.
118        """
119        from gameboy_worlds.interface.action import LowLevelAction, LowLevelActions
120
121        if transition_states is None:
122            return 0.0
123        current_state = transition_states[-1]
124        starter_chosen = current_state["pokemon_red_starter"]["current_starter"]
125        n_steps = current_state["core"]["steps"]
126        if starter_chosen is None:
127            if action == LowLevelAction and self.STEP_WISE_REWARD:
128                if "low_level_action" in action_kwargs:
129                    # reward for pressing A, penalty for pressing anything else
130                    low_level_action = action_kwargs["low_level_action"]
131                    if low_level_action == LowLevelActions.PRESS_BUTTON_A:
132                        return 0.5
133                    else:
134                        return -0.1
135            if n_steps >= self._emulator.max_steps - 2:  # some safety
136                return -1.0  # Penalty for not choosing a starter within max steps
137            else:
138                return 0.0
139        step_bonus = min(0.25 / (n_steps + 1), 0.5)
140        if starter_chosen == "charmander":
141            return 0.5 + step_bonus
142        else:
143            return (
144                0.25 + step_bonus
145            )  # Penalty for choosing the wrong starter. For now, just less reward.
146
147
148class PokemonRedChooseCharmanderEasyEnvironment(PokemonRedChooseCharmanderEnvironment):
149    def override_emulator_kwargs(emulator_kwargs: dict) -> dict:
150        """
151        Override default emulator keyword arguments for this environment.
152        """
153        Environment.override_state_tracker_class(
154            emulator_kwargs, PokemonRedStarterTracker
155        )
156        emulator_kwargs["init_state"] = "test_starter_easy"
157        return emulator_kwargs
158
159
160class PokemonRedChooseCharmanderHardEnvironment(PokemonRedChooseCharmanderEnvironment):
161    def override_emulator_kwargs(emulator_kwargs: dict) -> dict:
162        """
163        Override default emulator keyword arguments for this environment.
164        """
165        Environment.override_state_tracker_class(
166            emulator_kwargs, PokemonRedStarterTracker
167        )
168        emulator_kwargs["init_state"] = "test_starter_hard"
169        return emulator_kwargs