79 lines
2.8 KiB
Python
79 lines
2.8 KiB
Python
from __future__ import annotations
|
|
import random
|
|
import numpy as np
|
|
from typing import Callable
|
|
from retro.input import ProgrammaticInput
|
|
from retro.views.headless import HeadlessView
|
|
from retro_gamer.metadata import GameMetadata
|
|
from retro_gamer.observation import encode_observation
|
|
|
|
|
|
class GameEnvironment:
|
|
"""Gym-style wrapper around a retro Game for RL training.
|
|
|
|
Provides reset() / step(action) / observe(), managing one training episode
|
|
at a time. The game is restarted by calling the factory function on each reset.
|
|
"""
|
|
|
|
def __init__(self, game_factory: Callable, metadata: GameMetadata):
|
|
self.game_factory = game_factory
|
|
self.metadata = metadata
|
|
self.game = None
|
|
self.view: HeadlessView | None = None
|
|
self.inp: ProgrammaticInput | None = None
|
|
self._prev_reward: float = 0.0
|
|
|
|
def reset(self) -> np.ndarray:
|
|
"""Create a fresh game episode and return the initial observation."""
|
|
self.inp = ProgrammaticInput()
|
|
self.view = HeadlessView()
|
|
self.game = self.game_factory()
|
|
self.game.input_source = self.inp
|
|
self.game.view = self.view
|
|
self.game.start()
|
|
self._prev_reward = float(self.game.state.get(self.metadata.reward, 0))
|
|
return self._observe()
|
|
|
|
def step(self, action: str | None) -> tuple[np.ndarray, float, bool]:
|
|
"""Advance one turn with the given action.
|
|
|
|
action: a keystroke string (e.g. 'KEY_RIGHT') or None for no-op.
|
|
Returns (observation, reward, done).
|
|
Reward is the change in the reward state key since the previous step.
|
|
"""
|
|
self.inp.press(action)
|
|
self.game.step()
|
|
obs = self._observe()
|
|
reward = self._delta_reward()
|
|
done = not self.game.playing
|
|
return obs, reward, done
|
|
|
|
def _observe(self) -> np.ndarray:
|
|
return encode_observation(
|
|
self.view.board_characters,
|
|
dict(self.game.state),
|
|
self.metadata,
|
|
)
|
|
|
|
def _delta_reward(self) -> float:
|
|
current = float(self.game.state.get(self.metadata.reward, 0))
|
|
delta = current - self._prev_reward
|
|
self._prev_reward = current
|
|
return delta
|
|
|
|
def discover_character_set(self, exploration_turns: int) -> list[str]:
|
|
"""Run random turns to discover the characters that appear on the board.
|
|
Returns the sorted character list (excluding space).
|
|
"""
|
|
obs = self.reset()
|
|
chars: set[str] = set()
|
|
for _ in range(exploration_turns):
|
|
for row in self.view.board_characters:
|
|
chars.update(row)
|
|
action = random.choice(self.metadata.actions + [None])
|
|
_, _, done = self.step(action)
|
|
if done:
|
|
self.reset()
|
|
chars.discard(' ')
|
|
return sorted(chars)
|