Refactoring lab

games/babysnake/train.py

@@ -0,0 +1,41 @@
+"""Train a Q-learning agent to play BabySnake, then watch it play.
+
+Run this file to train and watch:
+
+    python train_babysnake.py
+
+This module wires the generic Q-learning algorithm in q_learning.py up to
+BabySnake specifically, using retro_gamer.GameEnvironment (configured by
+babysnake/pyproject.toml's observation_function) as the environment.
+"""
+
+from q_learning import QLearning
+from games import babysnake
+from retro.input import ProgrammaticInput
+from retro_gamer import GameEnvironment, GameMetadata
+
+def train():
+    trainer = QLearning()
+    env = GameEnvironment(babysnake.create_game, GameMetadata.from_pyproject("games.babysnake"))
+    return trainer.train(env, babysnake.ACTIONS)
+
+def watch(Q):
+    inp = ProgrammaticInput()
+
+    class PolicyInput:
+        """An input source that picks actions from the Q-table."""
+        def collect(self):
+            state = babysnake.get_state(game)
+            q, action = sorted([(Q.get((state, a), 0), a) for a in babysnake.ACTIONS], reverse=True)[0]
+            inp.press(action)
+            return inp.collect()
+
+    game = babysnake.create_game()
+    game.play(input_source=PolicyInput())
+
+if __name__ == '__main__':
+    print("Training Q-learning agent on BabySnake...")
+    Q = train()
+    print(f"\nDone. Q-table has {len(Q)} entries.")
+    print("\nWatching trained agent (press Enter or Escape to quit)...")
+    watch(Q)