gameboy_worlds.interface.pokemon.actions

  1from gameboy_worlds.utils import log_error, log_warn
  2from gameboy_worlds.interface.action import HighLevelAction, SingleHighLevelAction
  3from gameboy_worlds.emulation.pokemon.parsers import (
  4    AgentState,
  5    PokemonStateParser,
  6    BasePokemonRedStateParser,
  7)
  8from gameboy_worlds.emulation.pokemon.trackers import CorePokemonTracker
  9from gameboy_worlds.emulation import LowLevelActions
 10from abc import ABC, abstractmethod
 11from typing import List, Tuple, Dict
 12from gameboy_worlds.utils import show_frames
 13import numpy as np
 14
 15from gymnasium.spaces import Box, Discrete, Text, OneOf
 16import matplotlib.pyplot as plt
 17from PIL import Image
 18
 19HARD_MAX_STEPS = 5
 20""" The hard maximum number of steps we'll let agents take in a sequence """
 21
 22
 23def frame_changed(past: np.ndarray, preset: np.ndarray, epsilon=0.01):
 24    return np.abs(past - preset).mean() > epsilon
 25
 26
 27def _plot(past: np.ndarray, current: np.ndarray):
 28    # plot both side by side for debug
 29    fig, axs = plt.subplots(1, 2)
 30    axs[0].imshow(past)
 31    axs[1].imshow(current)
 32    plt.show()
 33
 34
 35class PassDialogueAction(SingleHighLevelAction):
 36    """
 37    Skips dialogue by pressing the B button.
 38
 39    Is Valid When:
 40    - In Dialogue State
 41
 42    Action Success Interpretation:
 43    - -1: Frame did not change
 44    - 0: Frame changed and no longer in dialogue state
 45    - 1: Frame changed but still in dialogue state
 46    """
 47
 48    REQUIRED_STATE_PARSER = PokemonStateParser
 49    REQUIRED_STATE_TRACKER = CorePokemonTracker
 50
 51    def is_valid(self, **kwargs):
 52        """
 53        Just checks if the agent is in dialogue state.
 54        """
 55        return (
 56            self._state_tracker.get_episode_metric(("pokemon_core", "agent_state"))
 57            == AgentState.IN_DIALOGUE
 58        )
 59
 60    def _execute(self):
 61        frames, done = self._emulator.step(LowLevelActions.PRESS_BUTTON_B)
 62        report = self._state_tracker.report()
 63        if not report["core"]["frame_changed"]:
 64            action_success = -1
 65        else:
 66            action_success = (
 67                0
 68                if self._emulator.state_parser.get_agent_state(frames[-1])
 69                != AgentState.IN_DIALOGUE
 70                else 1
 71            )
 72        return [report], action_success
 73
 74    @staticmethod
 75    def get_action_name() -> str:
 76        return "PassDialogue"
 77
 78
 79class InteractAction(SingleHighLevelAction):
 80    """
 81    Presses the A button to interact with an object in front of the agent.
 82
 83    Is Valid When:
 84    - In Free Roam State
 85
 86    Action Success Interpretation:
 87    - -1: Frame did not change or agent still in free roam state
 88    - 1: Agent not in free roam state
 89    """
 90
 91    REQUIRED_STATE_PARSER = PokemonStateParser
 92    REQUIRED_STATE_TRACKER = CorePokemonTracker
 93
 94    def is_valid(self, **kwargs):
 95        """
 96        Just checks if the agent is in free roam state.
 97        """
 98        return (
 99            self._state_tracker.get_episode_metric(("pokemon_core", "agent_state"))
100            == AgentState.FREE_ROAM
101        )
102
103    def _execute(self):
104        current_frame = self._emulator.get_current_frame()
105        frames, done = self._emulator.step(LowLevelActions.PRESS_BUTTON_A)
106        action_success = 0
107        # Check if the frames have changed. Be strict and require all to not permit jittering screens.
108        prev_frames = []
109        for frame in frames:
110            if (
111                self._emulator.state_parser.get_agent_state(frame)
112                != AgentState.FREE_ROAM
113            ):  # something happened lol
114                action_success = 1
115                break
116            for past_frame in prev_frames:
117                if not frame_changed(past_frame, frame):
118                    action_success = -1
119                    break
120            if action_success != 0:
121                break
122            prev_frames.append(frame)
123        if action_success == 0:
124            action_success = (
125                -1
126            )  # I guess? For some reason the previous thing doesn't catch same frames
127        return [
128            self._state_tracker.report()
129        ], action_success  # 0 means something maybe happened. 1 means def happened.
130
131    @staticmethod
132    def get_action_name() -> str:
133        return "Interact"
134
135
136class BaseMovementAction(HighLevelAction, ABC):
137    """
138    Base class for movement actions in the Pokemon environment.
139    Has utility methods for moving in directions.
140
141    Is Valid When:
142    - In Free Roam State
143
144    Action Success Interpretation:
145    - -1: Frame did not change, even on the first step
146    - 0: Finished all steps
147    - 1: Took some steps, but not all, and then frame stopped changing OR the frame starts oscillating (trying to check for jitter). This usually means we ran into an obstacle.
148    - 2: Took some steps, but agent state changed from free roam. This often means we entered a cutscene or battle.
149
150    Action Returns:
151    - `n_steps_taken` (`int`): Number of steps actually taken
152    - `rotated` (`bool` or `None`): True if the player has not moved, but has rotated. If the player has moved, this will be None. If it is False, it means the player tried to walk straight into an obstacle.
153
154    Known Limitations:
155    - Struggles to handle oscillating frames when the player is next to an obstacle. So if you are surrounded by water or bouncing flowers on all quadrants (e.g. the pier in Cinnabar Island), the system will think that the agent has moved forward, even though it has not.
156    """
157
158    REQUIRED_STATE_TRACKER = CorePokemonTracker
159    REQUIRED_STATE_PARSER = PokemonStateParser
160
161    def _is_uniform_quadrant(self, frame, quadrant_name) -> bool:
162        mapper = {
163            "screen_quadrant_1": "tr",
164            "screen_quadrant_2": "tl",
165            "screen_quadrant_3": "bl",
166            "screen_quadrant_4": "br",
167        }
168        quadrant_cells = self._emulator.state_parser.capture_grid_cells(
169            frame, quadrant=mapper[quadrant_name]
170        )
171        keys = list(quadrant_cells.keys())
172        x_min = 1e9
173        y_min = 1e9
174        x_max = -1e9
175        y_max = -1e9
176        for key in keys:
177            x_coord, y_coord = key
178            if x_coord * y_coord == 0:  # pop it to avoid checking player cell
179                quadrant_cells.pop(key)
180            else:
181                x_min = min(x_min, x_coord)
182                y_min = min(y_min, y_coord)
183                x_max = max(x_max, x_coord)
184                y_max = max(y_max, y_coord)
185        x_min = int(x_min)
186        y_min = int(y_min)
187        x_max = int(x_max)
188        y_max = int(y_max)
189        keys = list(quadrant_cells.keys())
190        vertical_uniform = True
191        horizontal_uniform = True
192        # check horizontal lines
193        for y in range(y_min + 1, y_max):  # avoid edges
194            first_cell = None
195            for x in range(x_min + 1, x_max):
196                cell = quadrant_cells[(x, y)]
197                if first_cell is None:
198                    first_cell = cell
199                else:
200                    if first_cell.shape == cell.shape:
201                        if frame_changed(first_cell, cell):
202                            horizontal_uniform = False
203                            break
204            if not horizontal_uniform:
205                break
206        if horizontal_uniform:
207            return True
208        # check vertical lines
209        for x in range(x_min + 1, x_max):
210            first_cell = None
211            for y in range(y_min + 1, y_max):
212                cell = quadrant_cells[(x, y)]
213                if first_cell is None:
214                    first_cell = cell
215                else:
216                    if first_cell.shape == cell.shape:
217                        if frame_changed(first_cell, cell):
218                            vertical_uniform = False
219                            break
220            if not vertical_uniform:
221                break
222        if vertical_uniform:
223            return True
224        return False
225
226    def judge_movement(
227        self, previous_frame: np.ndarray, current_frame: np.ndarray
228    ) -> Tuple[bool, bool]:
229        """
230        Judges whether movement has occurred between two frames.
231
232        Args:
233            previous_frame (np.ndarray): The previous frame.
234            current_frame (np.ndarray): The current frame.
235        Returns:
236            Tuple[bool, bool]: A tuple containing:
237            - bool: True if movement has occurred, False otherwise.
238            - bool: True if the player has not moved, but has rotated.
239        """
240        # if the full screen hasn't changed at all, player has neither moved nor rotated
241        if not frame_changed(previous_frame, current_frame):
242            return False, False
243        # split the screen into quadrants and check which quadrants have changed. If any of them stayed the same, the player has not moved, but may have rotated.
244        # One caveat is if the quadrant is uniform tiles (i.e. all have the same grid texture in them). In this case, we can't say for sure that the quadrant hasn't changed, since it may just be that the uniform texture is the same. So we check for that too.
245        flag = False
246        for quadrant in [
247            "screen_quadrant_1",
248            "screen_quadrant_2",
249            "screen_quadrant_3",
250            "screen_quadrant_4",
251        ]:
252            prev_quad = self._emulator.state_parser.capture_named_region(
253                previous_frame, quadrant
254            )
255            curr_quad = self._emulator.state_parser.capture_named_region(
256                current_frame, quadrant
257            )
258            prev_uniform = (
259                prev_quad.max() == prev_quad.min()  # screen is all black or all white
260                or self._is_uniform_quadrant(
261                    previous_frame, quadrant
262                )  # screen quadrant is uniform tiles
263            )
264            curr_uniform = (
265                curr_quad.max() == curr_quad.min()
266                or self._is_uniform_quadrant(current_frame, quadrant)
267            )
268            if (
269                not frame_changed(prev_quad, curr_quad)
270                and not prev_uniform
271                and not curr_uniform
272            ):  # then screen isn't just black, but also hasn't changed.
273                flag = True
274                break
275        if flag:  # then some frame stayed the same, so no movement, but maybe rotation.
276            prev_player_cell = self._emulator.state_parser.capture_grid_cells(
277                previous_frame
278            )[(0, 0)]
279            curr_player_cell = self._emulator.state_parser.capture_grid_cells(
280                current_frame
281            )[(0, 0)]
282            if frame_changed(prev_player_cell, curr_player_cell):
283                return False, True
284            else:
285                return False, False
286        else:
287            return True, None
288
289    def move(self, direction: str, steps: int) -> Tuple[np.ndarray, int]:
290        """
291        Move in a given direction for a number of steps.
292
293        :param direction: One of "up", "down", "left", "right"
294        :type direction: str
295        :param steps: Number of steps to move in that direction
296        :type steps: int
297        :return:  A tuple containing:
298
299                - A list of state tracker reports after each low level action executed. Length is equal to the number of low level actions executed.
300
301                - An integer action success status
302        :rtype: Tuple[ndarray[_AnyShape, dtype[Any]], int]
303        """
304        action_dict = {
305            "right": LowLevelActions.PRESS_ARROW_RIGHT,
306            "down": LowLevelActions.PRESS_ARROW_DOWN,
307            "up": LowLevelActions.PRESS_ARROW_UP,
308            "left": LowLevelActions.PRESS_ARROW_LEFT,
309        }
310        if direction not in action_dict.keys():
311            log_error(f"Got invalid direction to move {direction}", self._parameters)
312        action = action_dict[direction]
313        # keep trying the action.
314        # exit status 0 -> finished steps
315        # 1 -> took some steps, but not all, and then frame stopped changing OR the frame starts oscillating (trying to check for jitter)
316        # 2 -> took some steps, but agent state changed from free roam
317        # -1 -> frame didn't change, even on the first step
318        action_success = -1
319        transition_state_dicts = []
320        transition_frames = []
321        previous_frame = (
322            self._emulator.get_current_frame()
323        )  # Do NOT get the state tracker frame, as it may have a grid on it.
324        n_step = 0
325        n_successful_steps = 0
326        has_rotated = None
327        agent_state = AgentState.FREE_ROAM
328        while n_step < steps and agent_state == AgentState.FREE_ROAM:
329            frames, done = self._emulator.step(action)
330            transition_state_dicts.append(self._state_tracker.report())
331            transition_frames.extend(frames)
332            current_frame = (
333                self._emulator.get_current_frame()
334            )  # Do NOT use the emulator frame, as it may have a grid on it.
335            if done:
336                break
337            # check if frames changed. If not, break out.
338            # We check all frames in sequence to try and catch oscillations. But nothing will catch 1 step into wall in areas like this
339            player_moved, player_rotated = self.judge_movement(
340                previous_frame, current_frame
341            )
342            if player_rotated == True:
343                has_rotated = True
344            if not player_moved and not player_rotated:
345                break
346            if player_moved:
347                n_successful_steps += 1  # don't count rotation as a step
348            agent_state = self._emulator.state_parser.get_agent_state(
349                self._emulator.get_current_frame()
350            )
351            if agent_state != AgentState.FREE_ROAM:
352                break
353            n_step += 1
354            previous_frame = current_frame
355        if agent_state != AgentState.FREE_ROAM:
356            action_success = 2
357        else:
358            if n_step <= 0:
359                action_success = -1
360            elif n_step == steps:
361                action_success = 0
362            else:
363                action_success = 1
364        transition_state_dicts[-1]["core"]["action_return"] = {
365            "n_steps_taken": n_successful_steps,
366            "rotated": has_rotated,
367        }
368        return transition_state_dicts, action_success
369
370    def is_valid(self, **kwargs):
371        """
372        Just checks if the agent is in free roam state.
373        """
374        return (
375            self._state_tracker.get_episode_metric(("pokemon_core", "agent_state"))
376            == AgentState.FREE_ROAM
377        )
378
379
380class MoveStepsAction(BaseMovementAction):
381    """
382    Moves the agent in a specified cardinal direction for a specified number of steps.
383
384    Is Valid When:
385    - In Free Roam State
386    Action Success Interpretation:
387    - -1: Frame did not change, even on the first step
388    - 0: Finished all steps
389    - 1: Took some steps, but not all, and then frame stopped changing OR the frame starts oscillating (trying to check for jitter). This usually means we ran into an obstacle.
390    - 2: Took some steps, but agent state changed from free roam. This often means we entered a cutscene or battle.
391
392    Action Returns:
393    - `n_steps_taken` (`int`): Number of steps actually taken
394    - `rotated` (`bool` or `None`): True if the player has not moved, but has rotated. If the player has moved, this will be None. If it is False, it means the player tried to walk straight into an obstacle.
395    """
396
397    def get_action_space(self):
398        """
399        Returns a Box space representing movement in 2D.
400        The first dimension represents vertical movement (positive is up, negative is down).
401        The second dimension represents horizontal movement (positive is right, negative is left).
402
403        Returns:
404            Box: A Box space with shape (2,) and values ranging from -HARD_MAX_STEPS//2 to HARD_MAX_STEPS//2.
405
406        """
407        return Discrete(4 * HARD_MAX_STEPS)
408
409    def space_to_parameters(self, space_action):
410        direction = None
411        steps = None
412        if space_action < 0 or space_action >= 4 * HARD_MAX_STEPS:
413            # log_warn(f"Invalid space action {space_action}", self._parameters)
414            return None
415        if space_action < HARD_MAX_STEPS:
416            direction = "up"
417            steps = space_action
418        elif space_action < 2 * HARD_MAX_STEPS:
419            direction = "down"
420            steps = space_action - HARD_MAX_STEPS
421        elif space_action < 3 * HARD_MAX_STEPS:
422            direction = "left"
423            steps = space_action - 2 * HARD_MAX_STEPS
424        else:
425            direction = "right"
426            steps = space_action - 3 * HARD_MAX_STEPS
427        return {"direction": direction, "steps": steps + 1}
428
429    def parameters_to_space(self, direction: str, steps: int):
430        if steps <= 0 or steps > HARD_MAX_STEPS:
431            return None
432        if direction == "up":
433            return steps - 1
434        elif direction == "down":
435            return HARD_MAX_STEPS + steps - 1
436        elif direction == "left":
437            return 2 * HARD_MAX_STEPS + steps - 1
438        elif direction == "right":
439            return 3 * HARD_MAX_STEPS + steps - 1
440        else:
441            # log_warn(f"Unrecognized direction {direction}", self._parameters)
442            return None
443
444    def _execute(self, direction, steps):
445        transition_states, status = self.move(direction=direction, steps=steps)
446        return transition_states, status
447
448    def is_valid(self, **kwargs):
449        direction = kwargs.get("direction")
450        step = kwargs.get("step")
451        if direction is not None and direction not in ["up", "down", "left", "right"]:
452            return False
453        if step is not None:
454            if not isinstance(step, str):
455                return False
456            if step <= 0:
457                return False
458        return super().is_valid(**kwargs)
459
460    @staticmethod
461    def get_action_name(direction: str, steps: int) -> str:
462        return f"Move {direction} {steps}"
463
464
465class MoveGridAction(BaseMovementAction):
466    """
467    Moves the agent on both axes. Will always try to move right first and then up.
468
469    Is Valid When:
470    - In Free Roam State
471    Action Success Interpretation:
472    - -1: Frame did not change, even on the first step
473    - 0: Finished all steps
474    - 1: Took some steps, but not all, and then frame stopped changing OR the frame starts oscillating (trying to check for jitter). This usually means we ran into an obstacle.
475    - 2: Took some steps, but agent state changed from free roam. This often means we entered a cutscene or battle.
476
477    Action Returns:
478    - `n_steps_taken` (`int`): Number of steps actually taken
479    - `rotated` (`bool` or `None`): True if the player has not moved, but has rotated. If the player has moved, this will be None. If it is False, it means the player tried to walk straight into an obstacle.
480    """
481
482    def get_action_space(self):
483        """
484        Returns a Box space representing movement in 2D.
485        The first dimension represents vertical movement (positive is up, negative is down).
486        The second dimension represents horizontal movement (positive is right, negative is left).
487
488        Returns:
489            Box: A Box space with shape (2,) and values ranging from -HARD_MAX_STEPS//2 to HARD_MAX_STEPS//2.
490
491        """
492        return Box(
493            low=-HARD_MAX_STEPS // 2,
494            high=HARD_MAX_STEPS // 2,
495            shape=(2,),
496            dtype=np.int8,
497        )
498
499    def space_to_parameters(self, space_action):
500        right_action = space_action[0]
501        up_action = space_action[1]
502        return {"x_steps": right_action, "y_steps": up_action}
503
504    def parameters_to_space(self, x_steps, y_steps):
505        move_vec = np.zeros(2)  # x, y
506        move_vec[0] = x_steps
507        move_vec[1] = y_steps
508        return move_vec
509
510    def _execute(self, x_steps, y_steps):
511        x_direction = "right" if x_steps >= 0 else "left"
512        y_direction = "up" if y_steps >= 0 else "down"
513        if x_steps != 0:
514            transition_states, status = self.move(
515                direction=x_direction, steps=abs(x_steps)
516            )
517            if status != 0:
518                return transition_states, status
519        else:
520            transition_states = []
521        if y_steps != 0:
522            more_transition_states, status = self.move(
523                direction=y_direction, steps=abs(y_steps)
524            )
525            transition_states.extend(more_transition_states)
526        try:
527            status is not None
528        except NameError:
529            log_warn(
530                f"Weird case where both x_steps and y_steps are 0 in MoveGridAction or something. {x_steps}, {y_steps}",
531                self._parameters,
532            )
533            transition_states = [self._state_tracker.report()]
534            status = -1
535        return transition_states, status
536
537    def is_valid(self, x_steps: int = None, y_steps: int = None):
538        if x_steps is not None and y_steps is not None:
539            if not isinstance(x_steps, int) or not isinstance(y_steps, int):
540                return False
541            if x_steps == 0 and y_steps == 0:
542                return False
543        return super().is_valid()
544
545    @staticmethod
546    def get_action_name(x_steps: int, y_steps: int) -> str:
547        return f"MoveGrid ({x_steps}, {y_steps})"
548
549
550class MenuAction(HighLevelAction):
551    """
552    Allows simple navigation and option selection of menus.
553
554    Is Valid When:
555    - In Menu State
556
557    Action Success Interpretation:
558    - -1: Frame did not change.
559    - 0: Frame changed.
560    """
561
562    REQUIRED_STATE_PARSER = PokemonStateParser
563    REQUIRED_STATE_TRACKER = CorePokemonTracker
564
565    _MENU_ACTION_MAP = {
566        "up": LowLevelActions.PRESS_ARROW_UP,
567        "down": LowLevelActions.PRESS_ARROW_DOWN,
568        "confirm": LowLevelActions.PRESS_BUTTON_A,
569        "left": LowLevelActions.PRESS_ARROW_LEFT,
570        "right": LowLevelActions.PRESS_ARROW_RIGHT,
571        "back": LowLevelActions.PRESS_BUTTON_B,
572        # "open": LowLevelActions.PRESS_BUTTON_START,  # In general, we won't be using this action, but prefer OpenMenuAction instead.
573    }
574
575    _MENU_ACTION_KEYS = list(_MENU_ACTION_MAP.keys())
576
577    def is_valid(self, **kwargs):
578        """
579        Checks if the menu action is valid in the current state.
580
581        Args:
582            menu_action (str, optional): The menu action to check.
583        Returns:
584            bool: True if the action is valid, False otherwise.
585        """
586        menu_action = kwargs.get("menu_action", None)
587        if menu_action is not None:
588            if menu_action not in self._MENU_ACTION_KEYS:
589                return False
590        state = self._state_tracker.get_episode_metric(("pokemon_core", "agent_state"))
591        return state in [
592            AgentState.IN_MENU,
593            AgentState.IN_BATTLE,
594        ]  # works in battle screens and menu state
595
596    def get_action_space(self):
597        """
598        Returns a Discrete space representing menu actions.
599        Returns:
600            Discrete: A Discrete space with size equal to the number of menu actions.
601        """
602        return Discrete(len(self._MENU_ACTION_MAP))
603
604    def parameters_to_space(self, menu_action):
605        if menu_action not in self._MENU_ACTION_KEYS:
606            return None
607        return self._MENU_ACTION_KEYS.index(menu_action)
608
609    def space_to_parameters(self, space_action):
610        menu_action = None
611        if space_action < 0 or space_action >= len(self._MENU_ACTION_MAP):
612            return None
613        menu_action = self._MENU_ACTION_KEYS[space_action]
614        return {"menu_action": menu_action}
615
616    def _execute(self, menu_action):
617        action = self._MENU_ACTION_MAP[menu_action]
618        current_frame = self._emulator.get_current_frame()
619        frames, done = self._emulator.step(action)
620        action_success = 0 if frame_changed(current_frame, frames[-1]) else -1
621        return [self._state_tracker.report()], action_success
622
623    @staticmethod
624    def get_action_name(menu_action: str) -> str:
625        return f"Menu {menu_action}"
626
627
628class OpenMenuAction(HighLevelAction):
629    """
630    Opens the main menu from free roam and navigates to specified option.
631    Is Valid When:
632    - In Free Roam State
633
634    Action Success Interpretation:
635    - -1: Navigation Failure. Screen did not end up in expected end state. This is a bug and should not happen.
636    - 0: Navigation Success. Screen ended up in expected end state.
637
638    """
639
640    options = ["pokedex", "pokemon", "bag", "trainer"]
641
642    def get_action_space(self):
643        return Discrete(len(self.options))
644
645    def space_to_parameters(self, space_action):
646        if space_action < 0 or space_action >= len(self.options):
647            return None
648        return {"option": self.options[space_action]}
649
650    def parameters_to_space(self, option: str):
651        if option not in self.options:
652            return None
653        return self.options.index(option)
654
655    def is_valid(self, option: str = None):
656        if option is not None and option not in self.options:
657            return False
658        state = self._state_tracker.get_episode_metric(("pokemon_core", "agent_state"))
659        return state == AgentState.FREE_ROAM
660
661    def is_red_variant(self):
662        """
663        Returns true iff the current game is a red/blue variant.
664        """
665        return isinstance(self._emulator.state_parser, BasePokemonRedStateParser)
666
667    def _execute(self, option: str):
668        n_steps_down = 0
669        if option == "pokedex":
670            pass
671        elif option == "pokemon":
672            n_steps_down = 1
673        elif option == "bag":
674            n_steps_down = 2
675        elif option == "trainer":
676            if self.is_red_variant():
677                n_steps_down = 3
678            else:
679                n_steps_down = 4
680        else:
681            log_error(f"Invalid option {option}", self._parameters)
682        # first open menu
683        self._emulator.step(LowLevelActions.PRESS_BUTTON_START)
684        ret_states = [self._state_tracker.report()]
685        # go to the top
686        flag = False
687        for _ in range(6):
688            self._emulator.step(LowLevelActions.PRESS_ARROW_UP)
689            if self._emulator.state_parser.is_on_top_menu_option(
690                self._emulator.get_current_frame()
691            ):
692                flag = True
693                break
694        if not flag:  # could not get to top. Some error
695            return ret_states, -1
696        # go down n_steps_down
697        for _ in range(n_steps_down):
698            self._emulator.step(LowLevelActions.PRESS_ARROW_DOWN)
699            ret_states.append(self._state_tracker.report())
700        # confirm
701        self._emulator.step(LowLevelActions.PRESS_BUTTON_A)
702        ret_states.append(self._state_tracker.report())
703        return ret_states, 0
704
705    @staticmethod
706    def get_action_name(option: str) -> str:
707        return f"OpenMenu {option}"
708
709
710class BattleMenuAction(HighLevelAction):
711    """
712    Allows navigation of the battle menu.
713
714    Is Valid When:
715    - In Battle State
716
717    Action Success Interpretation:
718    - -1: Navigation Failure. Screen did not end up in expected end state.
719    - 0: Navigation Success. Screen ended up in expected end state. For run, got away safely.
720    - 1: Run Attempt Failed (cannot escape wild pokemon)
721    - 2: Run Attempt Failed (cannot escape trainer battle)
722    """
723
724    _OPTIONS = ["fight", "bag", "pokemon", "run", "progress"]
725
726    def is_valid(self, option: str = None):
727        if option is not None and option not in self._OPTIONS:
728            return False
729        state = self._state_tracker.get_episode_metric(("pokemon_core", "agent_state"))
730        # TODO: Must check we aren't in a 'learn new move' screen
731        return state == AgentState.IN_BATTLE
732
733    def get_action_space(self):
734        return Discrete(len(self._OPTIONS))
735
736    def get_all_valid_parameters(self):
737        state = self._state_tracker.get_episode_metric(("pokemon_core", "agent_state"))
738        if state != AgentState.IN_BATTLE:
739            return []
740        return [{"option": option} for option in self._OPTIONS]
741
742    def parameters_to_space(self, option: str):
743        return self._OPTIONS.index(option)
744
745    def space_to_parameters(self, space_action):
746        if space_action < 0 or space_action >= len(self._OPTIONS):
747            return None
748        return {"option": self._OPTIONS[space_action]}
749
750    def go_to_battle_menu(self):
751        # assumes we are in battle menu or will get there with some B's
752        state_reports = []
753        for i in range(3):
754            self._emulator.step(LowLevelActions.PRESS_BUTTON_B)
755            state_reports.append(self._state_tracker.report())
756        return state_reports
757
758    def button_sequence(self, low_level_actions: List[LowLevelActions]):
759        state_reports = self.go_to_battle_menu()
760        for action in low_level_actions:
761            self._emulator.step(action)
762        self._emulator.step(LowLevelActions.PRESS_BUTTON_A)  # confirm option
763        return state_reports + [self._state_tracker.report()]
764
765    def go_to_fight_menu(self):
766        return self.button_sequence(
767            [LowLevelActions.PRESS_ARROW_UP, LowLevelActions.PRESS_ARROW_LEFT]
768        )
769
770    def go_to_bag_menu(self):
771        return self.button_sequence(
772            [LowLevelActions.PRESS_ARROW_DOWN, LowLevelActions.PRESS_ARROW_LEFT]
773        )
774
775    def go_to_pokemon_menu(self):
776        return self.button_sequence(
777            [LowLevelActions.PRESS_ARROW_UP, LowLevelActions.PRESS_ARROW_RIGHT]
778        )
779
780    def go_to_run(self):
781        return self.button_sequence(
782            [LowLevelActions.PRESS_ARROW_DOWN, LowLevelActions.PRESS_ARROW_RIGHT]
783        )
784
785    def _execute(self, option):
786        success = -1
787        if option == "fight":
788            state_reports = self.go_to_fight_menu()
789            success = (
790                0
791                if self._emulator.state_parser.is_in_fight_options_menu(
792                    self._emulator.get_current_frame()
793                )
794                else -1
795            )
796        elif option == "bag":
797            state_reports = self.go_to_bag_menu()
798            success = (
799                0
800                if self._emulator.state_parser.is_in_fight_bag(
801                    self._emulator.get_current_frame()
802                )
803                else -1
804            )
805        elif option == "pokemon":
806            state_reports = self.go_to_pokemon_menu()
807            success = (
808                0
809                if self._emulator.state_parser.is_in_pokemon_menu(
810                    self._emulator.get_current_frame()
811                )
812                else -1
813            )
814        elif option == "run":
815            state_reports = self.go_to_run()
816            current_frame = self._emulator.get_current_frame()
817            got_away_safely = (
818                self._emulator.state_parser.named_region_matches_multi_target(
819                    current_frame, "dialogue_box_middle", "got_away_safely"
820                )
821            )
822            cannot_escape = (
823                self._emulator.state_parser.named_region_matches_multi_target(
824                    current_frame, "dialogue_box_middle", "cannot_escape"
825                )
826            )
827            cannot_run_from_trainer = (
828                self._emulator.state_parser.named_region_matches_multi_target(
829                    current_frame, "dialogue_box_middle", "cannot_run_from_trainer"
830                )
831            )
832            if got_away_safely:
833                success = 0
834                self._emulator.step(
835                    LowLevelActions.PRESS_BUTTON_B
836                )  # to clear the dialogue
837            elif cannot_escape:
838                success = 1
839                self._emulator.step(
840                    LowLevelActions.PRESS_BUTTON_B
841                )  # to clear the dialogue
842            elif cannot_run_from_trainer:
843                success = 2
844                self._emulator.step(LowLevelActions.PRESS_BUTTON_B)
845                state_reports.append(self._state_tracker.report())
846                self._emulator.step(
847                    LowLevelActions.PRESS_BUTTON_B
848                )  # Twice, to clear the dialogue
849            else:
850                pass  # Should never happen, but might.
851            state_reports.append(self._state_tracker.report())
852            return state_reports, success
853        elif option == "progress":
854            current_frame = self._emulator.get_current_frame()
855            state_reports = self.go_to_battle_menu()
856            new_frame = self._emulator.get_current_frame()
857            if frame_changed(current_frame, new_frame):
858                success = 0  # valid frame change, screen changed
859            else:
860                success = -1  # uneccesary progress press
861        else:
862            pass  # Will never happen.
863        return state_reports, success
864
865    @staticmethod
866    def get_action_name(option: str) -> str:
867        return f"BattleMenu {option}"
868
869
870class PickAttackAction(HighLevelAction):
871    """
872    Selects an attack option in the battle fight menu.
873
874    Is Valid When:
875    - In Battle State AND In Fight Menu
876
877    Action Success Interpretation:
878    - -1: Navigation Failure. Either could not get to the top of the attack menu (this should not happen) or the option index was too high (more likely the cause of failure).
879    - 0: Used attack successfully.
880    - 1: Tried to use a move with no PP remaining.
881    """
882
883    def get_action_space(self):
884        return Discrete(4)
885
886    def is_valid(self, option: int = None):
887        option = option - 1
888        if option is not None:
889            if option < 0 or option >= 4:
890                return False
891        return self._emulator.state_parser.is_in_fight_options_menu(
892            self._emulator.get_current_frame()
893        )
894
895    def parameters_to_space(self, option: int):
896        option = option - 1
897        return option
898
899    def space_to_parameters(self, space_action):
900        return {"option": space_action + 1}
901
902    def _execute(self, option: int):
903        # assume we are in the attack menu already
904        # first go to the top:
905        option = option - 1
906        flag = False
907        for _ in range(4):
908            self._emulator.step(LowLevelActions.PRESS_ARROW_UP)
909            if self._emulator.state_parser.is_on_top_attack_option(
910                self._emulator.get_current_frame()
911            ):
912                flag = True
913                break
914        if not flag:  # could not get to top. Some error
915            return [self._state_tracker.report()], -1
916        # then go down option times
917        for time in range(option):
918            self._emulator.step(LowLevelActions.PRESS_ARROW_DOWN)
919            if self._emulator.state_parser.is_on_top_attack_option(
920                self._emulator.get_current_frame()
921            ):
922                # went back to top, means that option was invalid
923                return [self._state_tracker.report()], -1
924        state_reports = []
925        self._emulator.step(LowLevelActions.PRESS_BUTTON_A)  # confirm option
926        state_reports.append(self._state_tracker.report())
927        if self._emulator.state_parser.tried_no_pp_move(
928            self._emulator.get_current_frame()
929        ):
930            self._emulator.step(
931                LowLevelActions.PRESS_BUTTON_B
932            )  # to clear the no PP dialogue
933            state_reports.append(self._state_tracker.report())
934            return state_reports, 1  # tried to use a move with no PP
935        else:
936            self._emulator.step(
937                LowLevelActions.PRESS_BUTTON_B
938            )  # to get through any attack animation dialogue
939            state_reports.append(self._state_tracker.report())
940        return state_reports, 0
941
942    @staticmethod
943    def get_action_name(option: int) -> str:
944        return f"PickAttack {option}"