Babysnake training works

train_babysnake.py

@@ -9,96 +9,33 @@ BabySnake specifically, using retro_gamer.GameEnvironment (configured by
 babysnake/pyproject.toml's observation_function) as the environment.
 """
 
+from q_learning import QLearning
 import babysnake
-import q_learning
-from babysnake_env import ACTIONS, get_state
 from retro.input import ProgrammaticInput
 from retro_gamer import GameEnvironment, GameMetadata
 
-
-def train(
-    episodes=1000,
-    alpha=0.1,
-    gamma=0.95,
-    epsilon=1.0,
-    epsilon_decay=0.995,
-    epsilon_min=0.05,
-    max_steps_per_episode=500,
-):
-    """Train a Q-learning agent on BabySnake.
-
-    Arguments:
-        episodes (int): How many episodes to run.
-        alpha (float): Learning rate.
-        gamma (float): Discount factor.
-        epsilon (float): Starting exploration rate.
-        epsilon_decay (float): Multiply epsilon by this each episode.
-        epsilon_min (float): Epsilon never falls below this.
-        max_steps_per_episode (int): Safety cutoff. Without this, a lucky
-            random walk that keeps finding food (each pickup restores more
-            energy than a turn costs) can make an episode run far longer
-            than intended, or even effectively forever.
-
-    Returns:
-        dict: The trained Q-table.
-    """
-    q_table = {}
+def train():
+    trainer = QLearning()
     env = GameEnvironment(babysnake.create_game, GameMetadata.from_pyproject("babysnake"))
+    return trainer.train(env, babysnake.ACTIONS)
 
-    for episode in range(episodes):
-        state = env.reset()
-        total_reward = 0.0
-
-        for _ in range(max_steps_per_episode):
-            if not env.game.playing:
-                break
-            action = q_learning.choose_action(q_table, state, ACTIONS, epsilon)
-            next_state, reward, done = env.step(action)
-            q_learning.update_q(q_table, state, action, reward, next_state, ACTIONS, alpha, gamma)
-            state = next_state
-            total_reward += reward
-
-        epsilon = max(epsilon_min, epsilon * epsilon_decay)
-
-        if (episode + 1) % 100 == 0:
-            print(
-                f"Episode {episode + 1:5d}  "
-                f"reward={total_reward:6.1f}  "
-                f"score={env.game.state['score']}  "
-                f"epsilon={epsilon:.3f}  "
-                f"q_entries={len(q_table)}"
-            )
-
-    return q_table
-
-
-def watch(q_table=None):
-    """Watch the trained agent play in the terminal.
-
-    Arguments:
-        q_table (dict | None): A trained Q-table. If None, trains first.
-    """
-    if q_table is None:
-        print("Training first...")
-        q_table = train()
-
+def watch(Q):
     inp = ProgrammaticInput()
 
     class PolicyInput:
         """An input source that picks actions from the Q-table."""
         def collect(self):
-            state = get_state(game)
-            action = q_learning.choose_action(q_table, state, ACTIONS, epsilon=0.0)
+            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_table = train()
-    print(f"\nDone. Q-table has {len(q_table)} entries.")
+    Q = train()
+    print(f"\nDone. Q-table has {len(Q)} entries.")
     print("\nWatching trained agent (press Enter or Escape to quit)...")
-    watch(q_table)
+    watch(Q)