0cd3c3b488
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.
208 lines
8.9 KiB
Python
208 lines
8.9 KiB
Python
from __future__ import annotations
|
|
import importlib
|
|
import tomllib
|
|
import tomli_w
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
|
|
@dataclass
|
|
class GameMetadata:
|
|
"""Describes a retro game for training purposes.
|
|
|
|
Required fields: actions, reward.
|
|
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
|
|
character_set: list[str] | None = None
|
|
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):
|
|
if not self.actions:
|
|
raise ValueError(
|
|
"The 'actions' list in [tool.retro-gamer] is empty or missing.\n"
|
|
"It should list the keyboard keys your agent can press, for example:\n\n"
|
|
' actions = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]\n\n'
|
|
"The agent will learn which actions lead to higher rewards."
|
|
)
|
|
if not isinstance(self.actions, list) or not all(isinstance(a, str) for a in self.actions):
|
|
raise ValueError(
|
|
f"'actions' must be a list of strings, but got: {self.actions!r}\n"
|
|
"Each entry should be a key name like \"KEY_RIGHT\" or \"KEY_SPACE\"."
|
|
)
|
|
if not self.reward:
|
|
raise ValueError(
|
|
"The 'reward' field in [tool.retro-gamer] is empty or missing.\n"
|
|
"It should name a game state variable whose value the agent is trying\n"
|
|
"to maximize — for example:\n\n"
|
|
" reward = \"score\"\n\n"
|
|
"The trainer watches how this value changes each step and uses those\n"
|
|
"changes as the reward signal."
|
|
)
|
|
if self.character_set is not None:
|
|
if not isinstance(self.character_set, list):
|
|
raise ValueError(
|
|
f"'character_set' must be a list of single characters, but got: {self.character_set!r}\n"
|
|
"Example: character_set = [\"@\", \"*\", \"#\"]"
|
|
)
|
|
for ch in self.character_set:
|
|
if not isinstance(ch, str) or len(ch) != 1:
|
|
raise ValueError(
|
|
f"Every entry in character_set must be a single character, but got {ch!r}.\n"
|
|
"Each character represents one type of cell on the game board.\n"
|
|
"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:
|
|
"""Load metadata from the [tool.retro-gamer] section of the game's pyproject.toml."""
|
|
pyproject_path = _find_pyproject(module_name)
|
|
if pyproject_path is None:
|
|
raise FileNotFoundError(
|
|
f"Could not find pyproject.toml for module '{module_name}'. "
|
|
f"Make sure the module is part of a Python project with a pyproject.toml."
|
|
)
|
|
with open(pyproject_path, 'rb') as f:
|
|
data = tomllib.load(f)
|
|
section = data.get('tool', {}).get('retro-gamer')
|
|
if section is None:
|
|
raise ValueError(
|
|
f"No [tool.retro-gamer] section found in {pyproject_path}.\n"
|
|
f"Add game metadata to your pyproject.toml:\n\n"
|
|
f"[tool.retro-gamer]\n"
|
|
f"actions = [\"KEY_RIGHT\", ...]\n"
|
|
f"reward = \"score\"\n"
|
|
)
|
|
return cls.from_dict(section)
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict) -> GameMetadata:
|
|
missing = [k for k in ('actions', 'reward') if k not in d]
|
|
if missing:
|
|
fields = ' and '.join(f"'{k}'" for k in missing)
|
|
raise ValueError(
|
|
f"The [tool.retro-gamer] section is missing required {fields}.\n"
|
|
"A minimal configuration looks like this:\n\n"
|
|
"[tool.retro-gamer]\n"
|
|
'actions = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]\n'
|
|
'reward = "score"\n\n'
|
|
"See the documentation for all available options."
|
|
)
|
|
board_size = tuple(d['board_size']) if 'board_size' in d else None
|
|
return cls(
|
|
actions=d['actions'],
|
|
reward=d['reward'],
|
|
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):
|
|
with open(path, 'wb') as f:
|
|
tomli_w.dump({'metadata': self.to_dict()}, f)
|
|
|
|
@property
|
|
def obs_size(self) -> int:
|
|
"""Total size of the flat observation vector."""
|
|
if not self.board:
|
|
return self.extras_size
|
|
C = len(self.character_set) if self.character_set else 0
|
|
bw, bh = self.board_size
|
|
return C * bw * bh + self.extras_size
|
|
|
|
@property
|
|
def n_actions(self) -> int:
|
|
"""Number of actions including no-op."""
|
|
return len(self.actions) + 1
|
|
|
|
|
|
def _find_pyproject(module_name: str) -> Path | None:
|
|
"""Walk up from a module's source file to find its pyproject.toml."""
|
|
try:
|
|
module = importlib.import_module(module_name)
|
|
except ImportError:
|
|
return None
|
|
module_file = getattr(module, '__file__', None)
|
|
if module_file is None:
|
|
return None
|
|
for parent in Path(module_file).resolve().parents:
|
|
candidate = parent / 'pyproject.toml'
|
|
if candidate.exists():
|
|
return candidate
|
|
if parent.name == 'site-packages':
|
|
break
|
|
return None
|