mwc/retro-gamer
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add observation_function for full custom control over the observation

Browse Source 
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.
This commit is contained in:
Chris Proctor
2026-06-23 20:45:48 -04:00
parent 426e59a54e
commit 0cd3c3b488
14 changed files with 479 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
retro_gamer/cli.py
View File

@@ -13,6 +13,12 @@ from retro_gamer.trainer import DQNTrainer, DEFAULTS, MODEL_KEYS
@click.group()
def cli():
"""Train and run RL agents for retro games."""
# Running the installed console script puts its own directory on sys.path,
# not the caller's cwd — but observation_function/game modules are
# typically plain files in the directory you ran retro-gamer from.
cwd = str(Path.cwd())
if cwd not in sys.path:
sys.path.insert(0, cwd)
# ---------------------------------------------------------------------------
@@ -79,8 +85,9 @@ def create(game, output, **hyperparams):
raise click.ClickException(str(e))
game_factory = _load_factory(game_config)
g = game_factory()
metadata.board_size = g.board_size
if metadata.board_size is None:
g = game_factory()
metadata.board_size = g.board_size
metadata.validate()
@@ -108,6 +115,8 @@ def create(game, output, **hyperparams):
else:
click.echo(f" characters : (will be auto-discovered during training)")
click.echo(f" architecture: {'CNN (spatial)' if metadata.spatial else 'MLP (non-spatial)'}")
if metadata.observation_function:
click.echo(f" observation : {metadata.observation_function} (custom)")
# ---------------------------------------------------------------------------

72
retro_gamer/env.py
View File

@@ -1,6 +1,5 @@
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
@@ -9,34 +8,49 @@ from retro_gamer.observation import encode_observation
class GameEnvironment:
"""Gym-style wrapper around a retro game for RL training."""
"""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,
egocentric: bool = False,
egocentric_player: str | None = None,
egocentric_radius: int | 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.egocentric = egocentric
self.egocentric_player = egocentric_player
self.egocentric_radius = egocentric_radius
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) -> np.ndarray:
"""Create a fresh game episode and return the initial observation."""
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()
@@ -46,7 +60,7 @@ class GameEnvironment:
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]:
def step(self, action: str | None) -> tuple:
"""Advance one turn. Returns (observation, reward, done)."""
self.inp.press(action)
self.game.step()
@@ -55,48 +69,18 @@ class GameEnvironment:
done = not self.game.playing
return obs, reward, done
def _observe(self) -> np.ndarray:
def _observe(self):
if self._observation_fn is not None:
return self._observation_fn(self.game)
state = dict(self.game.state)
if self.observe_state_sizes:
self._check_state_sizes(state)
player_pos = None
if self.egocentric and self.egocentric_player:
agent = self.game.get_agent_by_name(self.egocentric_player)
if agent is not None:
player_pos = agent.position
return encode_observation(
self.view.board_characters,
state,
self.metadata,
self.observe_state,
player_pos=player_pos,
egocentric_radius=self.egocentric_radius,
board=self.board,
)
def _check_state_sizes(self, state: dict):
for key, expected in self.observe_state_sizes.items():
val = state.get(key)
if val is None:
actual = 0
elif isinstance(val, (list, tuple)):
actual = len(val)
else:
actual = 1
if actual != expected:
raise ValueError(
f"State key '{key}' changed size during training:\n"
f" Expected : {expected} (discovered at training start)\n"
f" Got : {actual}\n\n"
f"This means game.state['{key}'] has a different length in some\n"
f"episodes than it had when training started. The neural network\n"
f"has a fixed input size and cannot adapt to changing state shapes.\n\n"
f"Fix: make sure create_game() always initializes '{key}' with a\n"
f"fixed-length value before the game starts each episode.\n"
f"For example, if '{key}' is a list of 9 values, it must always be\n"
f"a list of exactly 9 values — never more, never fewer, never missing."
)
def _delta_reward(self) -> float:
current = float(self.game.state.get(self.metadata.reward, 0))
delta = current - self._prev_reward

61
retro_gamer/metadata.py
View File

@@ -4,6 +4,7 @@ import tomllib
import tomli_w
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable
@dataclass
@@ -11,9 +12,15 @@ class GameMetadata:
"""Describes a retro game for training purposes.
Required fields: actions, reward.
Optional fields: character_set, spatial.
Discovered fields: board_size (from game.board_size), extras_size (from
the observe_state list in [preprocessing]).
Optional fields: character_set, spatial, observation_function.
Discovered fields: board_size (from game.board_size), extras_size
(computed by DQNTrainer from one sampled observation — never set in a
game's own pyproject.toml).
observation_function, if set, is a "module:attr" string naming a function
``f(game) -> Any`` that fully replaces the built-in board/observe_state
encoding. It is mutually exclusive with the [preprocessing] observe_state
option. See GameEnvironment for how the two paths are selected.
"""
actions: list[str]
reward: str
@@ -21,6 +28,7 @@ class GameMetadata:
spatial: bool = False
board: bool = True
board_size: tuple[int, int] | None = None
observation_function: str | None = None
extras_size: int = 0
def validate(self):
@@ -59,6 +67,48 @@ class GameMetadata:
"If you're not sure what characters your game uses, remove character_set\n"
"entirely and the trainer will discover them automatically."
)
if self.observation_function is not None:
if not isinstance(self.observation_function, str) or ':' not in self.observation_function:
raise ValueError(
f"'observation_function' must be a string of the form 'module:attr', "
f"but got: {self.observation_function!r}\n"
"Example: observation_function = \"my_game:get_observation\"\n"
"This should name a function f(game) -> observation that fully\n"
"describes what your agent observes each turn."
)
def resolve_observation_function(self) -> Callable | None:
"""Import and return the function named by observation_function, or None if unset.
Raises ValueError with an actionable message if the string isn't
"module:attr", the module can't be imported, or it has no such attribute.
"""
if self.observation_function is None:
return None
if ':' not in self.observation_function:
raise ValueError(
f"'observation_function' must be of the form 'module:attr', but got "
f"{self.observation_function!r} (no ':' found).\n"
"Example: observation_function = \"my_game:get_observation\""
)
module_name, attr_name = self.observation_function.split(':', 1)
try:
module = importlib.import_module(module_name)
except ImportError as e:
raise ValueError(
f"Could not import module {module_name!r} for observation_function "
f"{self.observation_function!r}: {e}\n"
"Make sure the module is importable (e.g. on PYTHONPATH or installed)."
) from e
try:
return getattr(module, attr_name)
except AttributeError:
raise ValueError(
f"Module {module_name!r} has no attribute {attr_name!r} "
f"(from observation_function = {self.observation_function!r}).\n"
f"Define a function named '{attr_name}' in {module_name} that takes a "
"game instance and returns its observation."
) from None
@classmethod
def from_pyproject(cls, module_name: str) -> GameMetadata:
@@ -102,17 +152,22 @@ class GameMetadata:
character_set=d.get('character_set'),
spatial=d.get('spatial', False),
board_size=board_size,
observation_function=d.get('observation_function'),
extras_size=d.get('extras_size', 0),
)
def to_dict(self) -> dict:
d = {
'actions': self.actions,
'reward': self.reward,
'extras_size': self.extras_size,
}
if self.board_size is not None:
d['board_size'] = list(self.board_size)
if self.character_set is not None:
d['character_set'] = self.character_set
if self.observation_function is not None:
d['observation_function'] = self.observation_function
return d
def to_toml(self, path: str | Path):

43
retro_gamer/model_agent.py
View File

@@ -45,17 +45,9 @@ class TrainedPolicy:
pre = config.get('preprocessing', {})
self._metadata.spatial = pre.get('spatial', False)
self._metadata.board = pre.get('board', True)
observe_state_sizes = pre.get('observe_state_sizes', {})
self._observe_state: list[str] = pre.get('observe_state', [])
self._egocentric: bool = pre.get('egocentric', False)
self._egocentric_player: str | None = pre.get('egocentric_player')
self._egocentric_radius: int | None = pre.get('egocentric_radius')
self._board: bool = pre.get('board', True)
if observe_state_sizes:
self._metadata.extras_size = sum(observe_state_sizes.values())
else:
self._metadata.extras_size = len(self._observe_state)
self._observation_fn = self._metadata.resolve_observation_function()
hyperparams = {**DEFAULTS, **config.get('model', {}), **config.get('training', {})}
self._model, _ = build_network(self._metadata, hyperparams)
@@ -78,26 +70,19 @@ class TrainedPolicy:
def get_action(self, game) -> str | None:
"""Return the key the model recommends this turn, or None for no-op."""
view = HeadlessView()
view.on_game_start(game)
view.render(game)
board_chars = view.board_characters
player_pos = None
if self._egocentric and self._egocentric_player:
agent = game.get_agent_by_name(self._egocentric_player)
if agent is not None:
player_pos = agent.position
obs = encode_observation(
board_chars,
dict(game.state),
self._metadata,
self._observe_state,
player_pos=player_pos,
egocentric_radius=self._egocentric_radius,
board=self._board,
)
if self._observation_fn is not None:
obs = self._observation_fn(game)
else:
view = HeadlessView()
view.on_game_start(game)
view.render(game)
obs = encode_observation(
view.board_characters,
dict(game.state),
self._metadata,
self._observe_state,
board=self._board,
)
device = next(self._model.parameters()).device
state_t = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0).to(device)

17
retro_gamer/observation.py
View File

@@ -64,25 +64,24 @@ def encode_observation(
state: dict,
metadata: GameMetadata,
observe_state: list[str],
player_pos: tuple[int, int] | None = None,
egocentric_radius: int | None = None,
board: bool = True,
) -> np.ndarray:
"""Encode board and/or selected state values into a flat 1D observation vector.
When *board* is True the board is encoded and prepended to the vector. If
player_pos and egocentric_radius are given the board is first cropped to a
(2r+1)×(2r+1) window centred on the player. For spatial games the board is
encoded channel-first (C, H, W) then flattened; for non-spatial games it is
encoded (H, W, C) then flattened. The state vector is appended at the end.
When *board* is True the board is encoded and prepended to the vector. For
spatial games the board is encoded channel-first (C, H, W) then flattened;
for non-spatial games it is encoded (H, W, C) then flattened. The state
vector is appended at the end.
When *board* is False only the observe_state features are returned.
For cropped (egocentric) boards, write a custom observation_function that
calls egocentric_board() and encode_board() directly instead of using this
function.
"""
if board:
if not metadata.character_set:
raise ValueError("character_set must be set before encoding observations")
if player_pos is not None and egocentric_radius is not None:
board_chars = egocentric_board(board_chars, player_pos, egocentric_radius)
board_enc = encode_board(board_chars, metadata.character_set) # (H, W, C)
if metadata.spatial:
board_vec = board_enc.transpose(2, 0, 1).flatten()

130
retro_gamer/trainer.py
View File

@@ -50,19 +50,17 @@ def _get_device() -> torch.device:
# Fields that make an existing checkpoint incompatible with the current config.
# Changing any of these requires starting training from scratch.
_INCOMPATIBLE_METADATA = {
'actions': 'the list of actions the agent can take (changes output layer size)',
'reward': 'the reward signal — Q-values trained on the old signal are meaningless for the new one',
'character_set': 'the set of board characters (changes input layer size)',
'board_size': 'the board dimensions (changes input layer size)',
'actions': 'the list of actions the agent can take (changes output layer size)',
'reward': 'the reward signal — Q-values trained on the old signal are meaningless for the new one',
'character_set': 'the set of board characters (changes input layer size)',
'board_size': 'the board dimensions (changes input layer size)',
'observation_function': 'how the observation is computed (changes input representation)',
'extras_size': 'the size of the non-board portion of the observation (changes input layer size)',
}
_INCOMPATIBLE_PREPROCESSING = {
'spatial': 'spatial vs non-spatial network type (changes network architecture)',
'board': 'whether the board is included in the observation (changes input size)',
'observe_state': 'the state keys included in the observation (changes input size)',
'observe_state_sizes': 'the size of each observed state key (changes input layer size)',
'egocentric': 'egocentric board transformation (changes input representation)',
'egocentric_player': 'the agent used as the egocentric center (changes input representation)',
'egocentric_radius': 'the egocentric crop radius (changes input layer size)',
}
_INCOMPATIBLE_ARCH = {
'hidden_sizes': 'the hidden layer sizes (changes network shape)',
@@ -274,11 +272,7 @@ class DQNTrainer:
pre = preprocessing or {}
self.observe_state: list[str] = pre.get('observe_state', [])
self.egocentric: bool = pre.get('egocentric', False)
self.egocentric_player: str | None = pre.get('egocentric_player', None)
self.egocentric_radius: int | None = pre.get('egocentric_radius', None)
self.board: bool = pre.get('board', True)
self.observe_state_sizes: dict[str, int] = pre.get('observe_state_sizes', {})
if self.board is False and metadata.spatial:
raise ValueError(
@@ -286,18 +280,11 @@ class DQNTrainer:
"A CNN requires a 2-D board to operate on. Either set spatial = false\n"
"or keep board = true."
)
if self.board is False and not self.observe_state:
if self.board is False and not self.observe_state and metadata.observation_function is None:
raise ValueError(
"preprocessing.board = false requires at least one entry in observe_state.\n"
"With board=false, the agent observes only the game state variables listed\n"
"in observe_state — if that list is empty, there is nothing to observe."
)
if self.egocentric and not self.egocentric_radius:
raise ValueError(
"preprocessing.egocentric = true requires egocentric_radius.\n"
"Choose a value based on how far the agent needs to see, e.g.:\n"
" egocentric_radius = 5 # 11×11 tight local view\n"
" egocentric_radius = 8 # 17×17 wider view"
"preprocessing.board = false requires at least one entry in observe_state\n"
"(or a metadata.observation_function). With board=false and no\n"
"observe_state, there is nothing for the agent to observe."
)
metadata.board = self.board
@@ -306,28 +293,16 @@ class DQNTrainer:
g = game_factory()
metadata.board_size = g.board_size
if self.egocentric_radius:
side = 2 * self.egocentric_radius + 1
metadata.board_size = (side, side)
self.env = GameEnvironment(
game_factory, metadata,
observe_state=self.observe_state,
egocentric=self.egocentric,
egocentric_player=self.egocentric_player,
egocentric_radius=self.egocentric_radius,
board=self.board,
observe_state_sizes=self.observe_state_sizes,
)
if metadata.character_set is None and self.board:
self._discover_character_set()
if self.observe_state and not self.observe_state_sizes:
self._discover_observe_state_sizes()
self.env.observe_state_sizes = self.observe_state_sizes
metadata.extras_size = sum(self.observe_state_sizes.values()) if self.observe_state_sizes else 0
self._discover_extras_size()
self.device = _get_device()
@@ -467,6 +442,7 @@ class DQNTrainer:
def _run_episode(self) -> tuple[float, int, float, bool]:
state = self.env.reset()
self._check_obs_length(state)
total_reward = 0.0
total_loss = 0.0
loss_count = 0
@@ -477,6 +453,7 @@ class DQNTrainer:
action_key = self._idx_to_key(action_idx)
next_state, reward, done = self.env.step(action_key)
self._check_obs_length(next_state)
self.memory.push(state, action_idx, reward, next_state, done)
if self.total_steps % self.hp['train_every'] == 0:
@@ -565,13 +542,9 @@ class DQNTrainer:
return {
'metadata': self.metadata.to_dict(),
'preprocessing': {
'spatial': self.metadata.spatial,
'board': self.board,
'observe_state': self.observe_state,
'observe_state_sizes': self.observe_state_sizes,
'egocentric': self.egocentric,
'egocentric_player': self.egocentric_player,
'egocentric_radius': self.egocentric_radius,
'spatial': self.metadata.spatial,
'board': self.board,
'observe_state': self.observe_state,
},
'hidden_sizes': self.hp['hidden_sizes'],
}
@@ -636,15 +609,61 @@ class DQNTrainer:
f"after {self.hp['exploration_turns']} exploration turns: {chars}"
)
def _discover_observe_state_sizes(self):
"""Sample game.state to determine the flat size of each observe_state key."""
self.env.reset()
state = dict(self.env.game.state)
sizes = {}
for key in self.observe_state:
val = state.get(key, 0)
sizes[key] = len(val) if isinstance(val, (list, tuple)) else 1
self.observe_state_sizes = sizes
def _discover_extras_size(self):
"""Sample one observation to determine extras_size (everything past the board).
This is also where the observation contract is enforced for training:
GameEnvironment itself has no opinion about what observations look like
(BabySnake, for example, uses it with a plain tuple), but DQNTrainer
needs a flat, numeric, fixed-length vector to feed a neural network.
"""
sample = self.env.reset()
try:
arr = np.asarray(sample, dtype=np.float32)
except (TypeError, ValueError) as e:
raise ValueError(
"Could not convert the observation to a numeric array for training:\n"
f" {sample!r}\n\n"
"DQNTrainer requires observation_function (or the built-in encoder) to\n"
"return a flat, numeric array-like value — a list/tuple of numbers or a\n"
"numpy array."
) from e
if arr.ndim != 1:
raise ValueError(
f"Expected a 1-D observation, but got shape {arr.shape}.\n"
"DQNTrainer always works with a single flat vector — board and extras\n"
"(if any) must be combined into one vector before being returned."
)
board_length = 0
if self.board:
C = len(self.metadata.character_set) if self.metadata.character_set else 0
bw, bh = self.metadata.board_size
board_length = C * bw * bh
if len(arr) < board_length:
raise ValueError(
f"The observation has length {len(arr)}, but character_set "
f"({C} chars) x board_size ({bw}x{bh}) = {board_length} is larger "
"than that.\n"
"Check that character_set/board_size match what your observation\n"
"actually encodes, or set board = false if there's no board in it."
)
self.metadata.extras_size = len(arr) - board_length
self._obs_len = len(arr)
def _check_obs_length(self, obs):
"""Raise a friendly error if an observation's length differs from the one discovered at init."""
length = len(np.asarray(obs, dtype=np.float32))
if length != self._obs_len:
raise ValueError(
"Observation length changed during training:\n"
f" Expected : {self._obs_len} (discovered at training start)\n"
f" Got : {length}\n\n"
"The neural network has a fixed input size and cannot adapt to a\n"
"changing observation shape. Make sure observation_function (or the\n"
"game's state) always produces the same length every episode."
)
def _save_config(self):
config_path = self.run_dir / 'config.toml'
@@ -658,13 +677,6 @@ class DQNTrainer:
pre['spatial'] = self.metadata.spatial
pre['board'] = self.board
pre['observe_state'] = self.observe_state
if self.observe_state_sizes:
pre['observe_state_sizes'] = self.observe_state_sizes
pre['egocentric'] = self.egocentric
if self.egocentric_player:
pre['egocentric_player'] = self.egocentric_player
if self.egocentric_radius:
pre['egocentric_radius'] = self.egocentric_radius
config['model'] = {k: v for k, v in self.hp.items() if k in MODEL_KEYS}
config['training'] = {k: v for k, v in self.hp.items() if k not in MODEL_KEYS}
with open(config_path, 'wb') as f: