42 lines
1.4 KiB
Python
42 lines
1.4 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.
|
|
"""
|
|
|
|
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)
|