Files
retro-gamer/retro_gamer/env.py
Chris Proctor 0cd3c3b488 Add observation_function for full custom control over the observation
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.
2026-06-23 20:45:48 -04:00

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)