gameboy_worlds.interface.controller

  1from abc import ABC
  2from typing import Any, Dict, Tuple, List, Optional, Type
  3from gameboy_worlds.utils import (
  4    verify_parameters,
  5    log_info,
  6    log_warn,
  7    log_error,
  8    load_parameters,
  9    get_lowest_level_subclass,
 10)
 11from gameboy_worlds.emulation.emulator import Emulator, LowLevelActions
 12from gameboy_worlds.interface.action import (
 13    HighLevelAction,
 14    LowLevelAction,
 15    RandomPlayAction,
 16    LowLevelPlayAction,
 17)
 18
 19import numpy as np
 20from gymnasium.spaces import OneOf, Space
 21
 22
 23class Controller(ABC):
 24    """
 25    Abstract base class for controllers interfacing with the emulator.
 26    Handles conversion between high level actions and Gym action spaces.
 27
 28    """
 29
 30    ACTIONS: List[Type[HighLevelAction]] = [HighLevelAction]
 31    """ A list of HighLevelAction classes that define the possible high level actions. 
 32    This is (almost) always, the only part that must be customized in subclasses.
 33    """
 34
 35    def __init__(self, parameters: Optional[dict] = None, seed: Optional[int] = None):
 36        self._parameters = load_parameters(parameters)
 37        self.actions: List[HighLevelAction] = [
 38            action(self._parameters) for action in self.ACTIONS
 39        ]
 40        """ A list of instantiated high level actions. """
 41        self.REQUIRED_STATE_TRACKER = get_lowest_level_subclass(
 42            [action.REQUIRED_STATE_TRACKER for action in self.actions]
 43        )
 44        """ The required state tracker class inferred from the high level actions. """
 45        self.action_space = OneOf(
 46            [action.get_action_space() for action in self.actions]
 47        )
 48        """ The Gym action Space consisting of a choice over all high level action spaces. """
 49        self.unassign_emulator()
 50        if seed is not None:
 51            self.seed(seed)
 52
 53    def seed(self, seed: Optional[int] = None):
 54        """
 55        Sets the random seed for the controller and its actions.
 56        Args:
 57            seed (int): The random seed to set.
 58        """
 59        self._rng = np.random.default_rng(seed)
 60        self.action_space.seed(seed)
 61        seed_value = seed
 62        for action in self.actions:
 63            if isinstance(seed, int):
 64                seed_value = (
 65                    seed + 1
 66                )  # Simple way to get different seeds for each action
 67            else:
 68                seed_value = None
 69            action.seed(seed_value)
 70
 71    def unassign_emulator(self):
 72        """
 73        Clears the reference to the emulator instance.
 74        """
 75        self._emulator = None
 76        for action in self.actions:
 77            action.unassign_emulator()
 78
 79    def assign_emulator(self, emulator: Emulator):
 80        """
 81        Sets a reference to the emulator instance.
 82        Args:
 83            emulator (Emulator): The emulator instance to be tracked.
 84        """
 85        for action in self.actions:
 86            action.assign_emulator(emulator)
 87        self._emulator = emulator
 88
 89    def get_action_space(self) -> OneOf:
 90        """
 91        Getter for the controller's Gym action space.
 92        Returns:
 93            OneOf: The Gym action Space consisting of a choice over all high level action spaces.
 94        """
 95        return self.action_space
 96
 97    def sample(self) -> OneOf:
 98        """
 99        Samples a random action from the controller's action space.
100        Returns:
101            OneOf: A random action from the controller's action space.
102        """
103        return self.action_space.sample()
104
105    def _space_action_to_high_level_action(
106        self, space_action: OneOf
107    ) -> Tuple[HighLevelAction, Dict[str, Any]]:
108        """
109        Interprets a Gym space action into a high level action and its parameters.
110
111        :param space_action: The action in the controller's action space.
112        :type space_action: OneOf
113        :return: The high level action and its parameters.
114        :rtype: Tuple[HighLevelAction, Dict[str, Any]]
115        """
116        action_index, space_action = space_action
117        action = self.actions[action_index]
118        action_class = self.ACTIONS[action_index]
119        parameters = action.space_to_parameters(space_action)
120        return action_class, parameters
121
122    def _high_level_action_to_space_action(
123        self, action: HighLevelAction, **kwargs
124    ) -> OneOf:
125        """
126        Converts a high level action and its parameters into a Gym Space action.
127
128        Args:
129            action (HighLevelAction): The high level action to convert.
130            **kwargs: Additional arguments required for the specific high level action.
131        Returns:
132            OneOf: The action in the controller's action space.
133        """
134        space_action = action.parameters_to_space(**kwargs)
135        if space_action is None:
136            return None
137        action_index = self.actions.index(action)
138        return (action_index, space_action)
139
140    def _emulator_running(self) -> bool:
141        """
142        Checks if the emulator is currently running.
143
144        Returns:
145            bool: True if the emulator is running, False otherwise.
146        """
147        if self._emulator is None:
148            log_error(
149                "Emulator reference not assigned to controller.", self._parameters
150            )
151        return not self._emulator.check_if_done()
152
153    def is_valid(self, action: Type[HighLevelAction], **kwargs) -> bool:
154        """
155        Checks if the specified high level action can be performed in the current state.
156
157        Args:
158            action (HighLevelAction): The high level action class to check.
159            **kwargs: Additional arguments required for the specific high level action.
160        Returns:
161            bool: True if the action is valid, False otherwise.
162        """
163        if not self._emulator_running():
164            return False
165        if action not in self.ACTIONS:
166            log_error(
167                "Action not recognized by controller. Are you passing in an instance of the action class?",
168                self._parameters,
169            )
170        # Find the action instance
171        action_index = self.ACTIONS.index(action)
172        checking_action = self.actions[action_index]
173        return checking_action.is_valid(**kwargs)
174
175    def get_valid_high_level_actions(
176        self,
177    ) -> Dict[Type[HighLevelAction], List[Dict[str, Any]]]:
178        """
179        Returns a list of all valid high level actions (including valid parameter inputs) that can be performed in the current state.
180
181        WARNING: Will fail if there are high level actions with infinite valid parameterizations. Use get_possibly_valid_high_level_actions() instead if that is the case.
182
183        :return: A dictionary mapping high level actions to their corresponding valid parameterizations.
184        :rtype: Dict[type[HighLevelAction], List[Dict[str, Any]]]
185        """
186        valid_actions = {}
187        if not self._emulator_running():
188            return valid_actions
189        for action in self.actions:
190            valid_parameters = action.get_all_valid_parameters()
191            if len(valid_parameters) > 0:
192                valid_actions[action] = valid_parameters
193        return valid_actions
194
195    def get_valid_space_actions(self) -> Dict[Type[HighLevelAction], List[OneOf]]:
196        """
197        Returns a list of valid actions in the controller's action space that can be performed in the current state.
198
199        WARNING: Will fail if there are high level actions with infinite valid parameterizations. Use get_possibly_valid_high_level_actions() instead if that is the case.
200
201        :return: A dictionary mapping high level actions to their corresponding valid space actions.
202        :rtype: Dict[type[HighLevelAction], List[OneOf]]
203        """
204        valid_space_actions = {}
205        if not self._emulator_running():
206            return valid_space_actions
207        valid_high_level_actions = self.get_valid_high_level_actions()
208        for action, parameter_list in valid_high_level_actions.items():
209            valid_space_actions[action] = []
210            for parameters in parameter_list:
211                space_action = self._high_level_action_to_space_action(
212                    action, **parameters
213                )
214                if space_action is None:
215                    log_error(
216                        f"Invalid action parameters combination for {action}: {parameters}. Ensure there are no bugs in {action}.get_all_valid_parameters",
217                        self._parameters,
218                    )
219                valid_space_actions[action].append(space_action)
220        return valid_space_actions
221
222    def get_possibly_valid_high_level_actions(self) -> List[Type[HighLevelAction]]:
223        """
224        Returns a list of valid high level actions that can be performed (with some parameterized input) in the current state.
225
226        Returns:
227            List[Type[HighLevelAction]]: A list of valid high level actions.
228        """
229        if not self._emulator_running():
230            return []
231        actions = []
232        for i, action_class in enumerate(self.ACTIONS):
233            action = self.actions[i]
234            if action.is_valid():
235                actions.append(action_class)
236        return actions
237
238    def execute_space_action(
239        self, action: OneOf
240    ) -> Tuple[Optional[List[Dict[str, Dict[str, Any]]]], Optional[int]]:
241        """
242        Executes the specified high level action on the emulator after checking for validity.
243
244        :param action: The action in the controller's action space.
245        :type action: OneOf
246        :return:
247            - List[Dict[str, Dict[str, Any]]]: A list of state tracker reports after each low level action executed. Length is equal to the number of low level actions executed.
248
249            - int: Action success status.
250        :rtype: Tuple[List[Dict[str, Dict[str, Any]]] | None, int | None]
251        """
252        action_index, space_action = action
253        executing_action: HighLevelAction = self.actions[action_index]
254        return executing_action.execute_space_action(space_action)
255
256    def execute(
257        self, action: Type[HighLevelAction], **kwargs
258    ) -> Tuple[Optional[List[Dict[str, Dict[str, Any]]]], Optional[int]]:
259        """
260        Executes the specified high level action on the emulator after checking for validity.
261
262        :param action: The HighLevelAction
263        :type action: Type[HighLevelAction]
264        :param kwargs: Additional arguments required for the specific high level action.
265        :type kwargs: Dict[str, Any]
266        :return:
267            - List[Dict[str, Dict[str, Any]]]: A list of state tracker reports after each low level action executed. Length is equal to the number of low level actions executed.
268
269            - int: Action success status.
270        :rtype: Tuple[List[Dict[str, Dict[str, Any]]] | None, int | None]
271        """
272        if action not in self.ACTIONS:
273            log_error(
274                "Action not recognized by controller. Are you passing in an instance of the action class?",
275                self._parameters,
276            )
277        # Find the action instance
278        action_index = self.ACTIONS.index(action)
279        executing_action = self.actions[action_index]
280        return executing_action.execute(**kwargs)
281
282    def string_to_space_action(self, input_str: str) -> Optional[OneOf]:
283        """
284        Converts a string input to a space action
285        Args:
286            input_str (str): The string input representing the high level action and its parameters.
287
288        Returns:
289            OneOf: The action in the controller's action space.
290        """
291        action, kwargs = self.string_to_high_level_action(input_str=input_str)
292        if action is None:
293            return None
294        return self._high_level_action_to_space_action(action, kwargs)
295
296    def execute_string(
297        self, input_str: str
298    ) -> Tuple[Optional[List[Dict[str, Dict[str, Any]]]], Optional[int]]:
299        """
300        Executes the high level action implied by the input string.
301
302        :param input_str: String representing the high level action and its parameters.
303        :type input_str: str
304        :param kwargs: Additional arguments required for the specific high level action.
305        :type kwargs: Dict[str, Any]
306        :return:
307            - List[Dict[str, Dict[str, Any]]]: A list of state tracker reports after each low level action executed. Length is equal to the number of low level actions executed. Is None if the input string does not map to a valid action.
308
309            - int: Action success status. Is None if the input string does not map to a valid action.
310        :rtype: Tuple[List[Dict[str, Dict[str, Any]]] | None, int | None]
311        """
312        action, kwargs = self.string_to_high_level_action(input_str=input_str)
313        if action is None:
314            return None, None
315        return self.execute(action, kwargs)
316
317    def string_to_high_level_action(
318        self, input_str: str
319    ) -> Tuple[Type[HighLevelAction], Dict[str, Any]]:
320        """
321        Provide a way to map a string input to a HighLevelAction and parameters.
322
323        Implement if you want to use the human_step_play method, or if you want to allow a LM based agent to give its actions in text.
324        Must return None, None if the input_str does not map to an action.
325        """
326        raise NotImplementedError
327
328    def get_action_strings(
329        self, return_all: bool = False
330    ) -> Dict[HighLevelAction, str]:
331        """
332        Provide a way to verbalize the allowed high level actions, along with the format of the input parameters.
333        Useful for prompting a VLM to choose an action.
334
335        This should match the mapping in string_to_high_level_action
336
337        :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.
338        :type return_all: bool
339        :return: A dictionary mapping high level actions to their verbalizations and input formats.
340        :rtype: Dict[HighLevelAction, str]
341        """
342        raise NotImplementedError
343
344
345def parse_button_string(input_str: str) -> Optional[LowLevelActions]:
346    string_low = input_str.lower().strip().strip("(").strip(")").replace("_", " ")
347    mapper = {
348        "a": LowLevelActions.PRESS_BUTTON_A,
349        "up": LowLevelActions.PRESS_ARROW_UP,
350        "b": LowLevelActions.PRESS_BUTTON_B,
351        "down": LowLevelActions.PRESS_ARROW_DOWN,
352        "left": LowLevelActions.PRESS_ARROW_LEFT,
353        "right": LowLevelActions.PRESS_ARROW_RIGHT,
354        "start": LowLevelActions.PRESS_BUTTON_START,
355        "u": LowLevelActions.PRESS_ARROW_UP,
356        "d": LowLevelActions.PRESS_ARROW_DOWN,
357        "l": LowLevelActions.PRESS_ARROW_LEFT,
358        "r": LowLevelActions.PRESS_ARROW_RIGHT,
359        "s": LowLevelActions.PRESS_BUTTON_START,
360    }
361    for map_opt in mapper:
362        if string_low == map_opt:
363            return mapper[map_opt]
364    if "move" in string_low:
365        directions_in_string = []
366        for dir_opt in ["up", "down", "left", "right"]:
367            if dir_opt in string_low:
368                directions_in_string.append(dir_opt)
369        if len(directions_in_string) == 1:
370            return mapper[directions_in_string[0]]
371    if "press" in string_low:
372        buttons_in_string = []
373        for button_opt in ["a", "b", "start"]:
374            if button_opt in string_low:
375                buttons_in_string.append(button_opt)
376        if len(buttons_in_string) == 1:
377            return mapper[buttons_in_string[0]]
378    return None
379
380
381class LowLevelController(Controller):
382    """A controller that executes low level actions directly on the emulator."""
383
384    ACTIONS = [LowLevelAction]
385    """ A HighLevelAction subclass that directly maps to low level actions. """
386
387    def string_to_high_level_action(self, input_str):
388        low_level_action = parse_button_string(input_str)
389        if low_level_action is None:
390            return None, None
391        return LowLevelAction, {"low_level_action": low_level_action}
392
393    def get_action_strings(self, return_all=False):
394        msg = f"""
395        Arrow Keys (UP for up, DOWN for down, LEFT for left, RIGHT for right), A and B for buttons, START for start.
396        """
397        return {LowLevelAction: msg}
398
399
400class LowLevelPlayController(Controller):
401    """A controller that executes low level actions directly, but no menu button presses."""
402
403    ACTIONS = [LowLevelPlayAction]
404    """ A HighLevelAction subclass that directly maps to low level actions, but no menu button presses. """
405
406    def string_to_high_level_action(self, input_str):
407        low_level_action = parse_button_string(input_str)
408        if low_level_action is None:
409            return None, None
410        return LowLevelPlayAction, {"low_level_action": low_level_action}
411
412    def get_action_strings(self):
413        msg = f"""
414        A, B for button. LEFT, RIGHT, UP, DOWN for arrow keys
415        """
416        return msg
417
418
419class RandomPlayController(Controller):
420    """A controller that performs random play on the emulator using low level actions."""
421
422    ACTIONS = [RandomPlayAction]
423    """ A HighLevelAction subclass that performs random low level actions. """
424
425
426_ALWAYS_VALID_CONTROLLERS = {
427    "low_level": LowLevelController,
428    "low_level_play": LowLevelPlayController,
429    "random_play": RandomPlayController,
430}
431""" Controllers that are always valid for any game and environment. """