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.

retro_gamer/env.py

@@ -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