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.
This commit is contained in:
Chris Proctor
208
tests/test_observation_function.py
Normal file
208
tests/test_observation_function.py
Normal file
@@ -0,0 +1,208 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user