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/trainer.py

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