Lets a game define how its state becomes an observation as a plain Python function (module:attr), used identically by GameEnvironment (training) and TrainedPolicy (inference) instead of two independently maintained encoding paths. Removes the egocentric/egocentric_player/ egocentric_radius flags — cropping is now something an observation_function does itself by calling egocentric_board(), and extras_size is discovered from one sampled observation instead of being configured via observe_state_sizes.
103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
from __future__ import annotations
|
|
import random
|
|
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.
|
|
|
|
The observation returned by reset()/step() comes from one of two mutually
|
|
exclusive paths: metadata.observation_function, if set, fully replaces the
|
|
built-in board/observe_state encoding — see GameMetadata for its contract.
|
|
Otherwise the built-in encoder (encode_observation) is used, configured by
|
|
observe_state/board/observe_state_sizes below.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
game_factory: Callable,
|
|
metadata: GameMetadata,
|
|
observe_state: list[str] | None = None,
|
|
board: bool = True,
|
|
observe_state_sizes: dict[str, int] | None = None,
|
|
):
|
|
self.game_factory = game_factory
|
|
self.metadata = metadata
|
|
self.observe_state = observe_state or []
|
|
self.board = board
|
|
self.observe_state_sizes = observe_state_sizes or {}
|
|
self._observation_fn = metadata.resolve_observation_function()
|
|
if self._observation_fn is not None and self.observe_state:
|
|
raise ValueError(
|
|
"Both metadata.observation_function and [preprocessing].observe_state "
|
|
"are set, but they're two conflicting ways of describing the\n"
|
|
"observation. Use observation_function for full custom control, or\n"
|
|
"observe_state (with the built-in board encoder) but not both."
|
|
)
|
|
self.game = None
|
|
self.view: HeadlessView | None = None
|
|
self.inp: ProgrammaticInput | None = None
|
|
self._prev_reward: float = 0.0
|
|
|
|
def reset(self):
|
|
"""Create a fresh game episode and return the initial observation.
|
|
|
|
The observation's type depends on metadata.observation_function: a
|
|
numpy array when using the built-in encoder (or a custom function
|
|
built for DQN training), but it can be anything a custom function
|
|
returns — e.g. a plain tuple for tabular use.
|
|
"""
|
|
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:
|
|
"""Advance one turn. Returns (observation, reward, done)."""
|
|
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):
|
|
if self._observation_fn is not None:
|
|
return self._observation_fn(self.game)
|
|
state = dict(self.game.state)
|
|
return encode_observation(
|
|
self.view.board_characters,
|
|
state,
|
|
self.metadata,
|
|
self.observe_state,
|
|
board=self.board,
|
|
)
|
|
|
|
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."""
|
|
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)
|