e8a24ae7be
q_learning.py mixed a bespoke BabySnake environment wrapper, the two functions students implement, a training loop, and a terminal watch routine in one file, with no tests and no single command to run training. - q_learning.py / q_learning_solution.py now hold only choose_action and update_q, with actions as an explicit parameter instead of a module-global, so the file has no babysnake/retro/retro_gamer imports at all. - train_babysnake.py builds GameEnvironment directly from babysnake's pyproject metadata, trains, and watches the trained agent in one command: `python train_babysnake.py`. It also caps steps per episode, since a lucky random walk that keeps finding food can otherwise make an episode run unboundedly long. - test_q_learning.py adds unittest coverage for both functions with no game dependency. - questions.md adds a checkpoint instructing students to get the tests passing before training, and points the post-training step at train_babysnake.py.
105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
"""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.
|
|
"""
|
|
|
|
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 = {}
|
|
env = GameEnvironment(babysnake.create_game, GameMetadata.from_pyproject("babysnake"))
|
|
|
|
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()
|
|
|
|
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)
|
|
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.")
|
|
print("\nWatching trained agent (press Enter or Escape to quit)...")
|
|
watch(q_table)
|