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.
209 lines
7.3 KiB
Python
209 lines
7.3 KiB
Python
"""Tests for the observation_function machinery: GameMetadata resolution,
|
|
GameEnvironment delegation, and DQNTrainer's extras_size discovery.
|
|
|
|
Run with: uv run python -m unittest tests/test_observation_function.py -v
|
|
"""
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest import TestCase, main
|
|
|
|
import numpy as np
|
|
from retro.game import Game
|
|
|
|
from retro_gamer.metadata import GameMetadata
|
|
from retro_gamer.observation import encode_observation
|
|
from retro_gamer.env import GameEnvironment
|
|
from retro_gamer.trainer import DQNTrainer
|
|
|
|
ACTIONS = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]
|
|
|
|
|
|
def sample_observation_fn(game):
|
|
"""A trivial observation_function used to test dotted-path resolution."""
|
|
return "observation-from-sample_observation_fn"
|
|
|
|
|
|
def tuple_observation_fn(game):
|
|
"""Returns a plain tuple — exercises the non-array tabular use case."""
|
|
return (1, 2, 3)
|
|
|
|
|
|
def make_game_factory(board_size=(2, 2)):
|
|
def factory():
|
|
return Game([], {'reward': 0.0, 'score': 0}, board_size=board_size, show_state=False)
|
|
return factory
|
|
|
|
|
|
def obs_with_extras(game):
|
|
"""Board (character_set=2, board_size=(2,2) -> 8) + 3 extras = 11."""
|
|
return np.zeros(2 * 2 * 2 + 3, dtype=np.float32)
|
|
|
|
|
|
def obs_2d(game):
|
|
"""Not 1-D — used to test the training-time shape validation."""
|
|
return np.zeros((2, 2), dtype=np.float32)
|
|
|
|
|
|
def obs_too_short(game):
|
|
"""Shorter than character_set (2) x board_size (2x2) = 8."""
|
|
return np.zeros(5, dtype=np.float32)
|
|
|
|
|
|
_changing_obs_counter = {"n": 0}
|
|
|
|
|
|
def changing_obs(game):
|
|
"""Returns length 8 on the first call, length 9 on every call after."""
|
|
_changing_obs_counter["n"] += 1
|
|
length = 8 if _changing_obs_counter["n"] == 1 else 9
|
|
return np.zeros(length, dtype=np.float32)
|
|
|
|
|
|
class TestResolveObservationFunction(TestCase):
|
|
def test_returns_none_when_unset(self):
|
|
metadata = GameMetadata(actions=ACTIONS, reward="reward")
|
|
self.assertIsNone(metadata.resolve_observation_function())
|
|
|
|
def test_resolves_valid_dotted_path(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
observation_function=f"{__name__}:sample_observation_fn",
|
|
)
|
|
fn = metadata.resolve_observation_function()
|
|
self.assertIs(fn, sample_observation_fn)
|
|
|
|
def test_raises_on_malformed_string(self):
|
|
metadata = GameMetadata(actions=ACTIONS, reward="reward", observation_function="no_colon_here")
|
|
with self.assertRaises(ValueError):
|
|
metadata.resolve_observation_function()
|
|
|
|
def test_raises_on_unimportable_module(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
observation_function="this_module_does_not_exist:fn",
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
metadata.resolve_observation_function()
|
|
|
|
def test_raises_on_missing_attribute(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
observation_function=f"{__name__}:no_such_function",
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
metadata.resolve_observation_function()
|
|
|
|
def test_validate_raises_on_malformed_observation_function(self):
|
|
metadata = GameMetadata(actions=ACTIONS, reward="reward", observation_function="no_colon_here")
|
|
with self.assertRaises(ValueError):
|
|
metadata.validate()
|
|
|
|
def test_validate_passes_with_well_formed_observation_function(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
observation_function=f"{__name__}:sample_observation_fn",
|
|
)
|
|
metadata.validate() # should not raise
|
|
|
|
|
|
class TestEncodeObservation(TestCase):
|
|
def setUp(self):
|
|
self.metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
character_set=["@", "*"], board_size=(2, 2),
|
|
)
|
|
self.board_chars = [["@", " "], [" ", "*"]]
|
|
|
|
def test_board_plus_extras(self):
|
|
obs = encode_observation(
|
|
self.board_chars, {"x": 1.0, "y": 2.0}, self.metadata, ["x", "y"],
|
|
)
|
|
self.assertEqual(obs.shape, (2 * 2 * 2 + 2,))
|
|
|
|
def test_board_false_returns_only_extras(self):
|
|
obs = encode_observation(
|
|
self.board_chars, {"x": 1.0, "y": 2.0}, self.metadata, ["x", "y"], board=False,
|
|
)
|
|
np.testing.assert_array_equal(obs, [1.0, 2.0])
|
|
|
|
|
|
class TestGameEnvironment(TestCase):
|
|
def test_custom_observation_function_returns_its_value_directly(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
observation_function=f"{__name__}:tuple_observation_fn",
|
|
)
|
|
env = GameEnvironment(make_game_factory(), metadata)
|
|
self.assertEqual(env.reset(), (1, 2, 3))
|
|
obs, reward, done = env.step(None)
|
|
self.assertEqual(obs, (1, 2, 3))
|
|
|
|
def test_mutual_exclusivity_raises(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
observation_function=f"{__name__}:tuple_observation_fn",
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
GameEnvironment(make_game_factory(), metadata, observe_state=["score"])
|
|
|
|
def test_built_in_path_still_works(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
character_set=["@", "*"], board_size=(2, 2),
|
|
)
|
|
env = GameEnvironment(make_game_factory(), metadata, observe_state=["score"])
|
|
obs = env.reset()
|
|
self.assertEqual(obs.shape, (2 * 2 * 2 + 1,))
|
|
|
|
|
|
class TestDQNTrainerExtrasSize(TestCase):
|
|
def _trainer(self, metadata, preprocessing=None, board_size=(2, 2)):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(tmp.cleanup)
|
|
return DQNTrainer(
|
|
make_game_factory(board_size), metadata, Path(tmp.name),
|
|
preprocessing=preprocessing,
|
|
)
|
|
|
|
def test_computes_extras_size_from_custom_function(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
character_set=["@", "*"], board_size=(2, 2),
|
|
observation_function=f"{__name__}:obs_with_extras",
|
|
)
|
|
trainer = self._trainer(metadata)
|
|
self.assertEqual(trainer.metadata.extras_size, 3)
|
|
|
|
def test_raises_on_non_1d_observation(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
character_set=["@", "*"], board_size=(2, 2),
|
|
observation_function=f"{__name__}:obs_2d",
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
self._trainer(metadata)
|
|
|
|
def test_raises_when_observation_too_short_for_declared_board(self):
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
character_set=["@", "*"], board_size=(2, 2),
|
|
observation_function=f"{__name__}:obs_too_short",
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
self._trainer(metadata)
|
|
|
|
def test_raises_when_observation_length_changes_mid_run(self):
|
|
_changing_obs_counter["n"] = 0
|
|
metadata = GameMetadata(
|
|
actions=ACTIONS, reward="reward",
|
|
observation_function=f"{__name__}:changing_obs",
|
|
)
|
|
trainer = self._trainer(metadata, preprocessing={'board': False})
|
|
self.assertEqual(trainer.metadata.extras_size, 8)
|
|
with self.assertRaises(ValueError):
|
|
trainer._run_episode()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|