diff --git a/babysnake/__init__.py b/babysnake/__init__.py index 776d2cc..25fcb77 100644 --- a/babysnake/__init__.py +++ b/babysnake/__init__.py @@ -9,6 +9,14 @@ The game ends when energy reaches 0. from random import randint, choice from retro.game import Game +ACTIONS = ["KEY_RIGHT", "KEY_DOWN", "KEY_LEFT", "KEY_UP"] + + +def get_state(game): + """Return the game state as a hashable (agent_x, agent_y, food_x, food_y) tuple.""" + s = game.state + return (int(s['agent_x']), int(s['agent_y']), int(s['food_x']), int(s['food_y'])) + BOARD_SIZE = 4 START_ENERGY = 50 FOOD_ENERGY = 30 diff --git a/babysnake/pyproject.toml b/babysnake/pyproject.toml index c605c41..6b68000 100644 --- a/babysnake/pyproject.toml +++ b/babysnake/pyproject.toml @@ -1,4 +1,4 @@ [tool.retro-gamer] actions = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"] reward = "reward" -observation_function = "babysnake_env:get_state" +observation_function = "babysnake:get_state" diff --git a/babysnake_env.py b/babysnake_env.py deleted file mode 100644 index 9f5d2f5..0000000 --- a/babysnake_env.py +++ /dev/null @@ -1,13 +0,0 @@ -"""BabySnake's observation_function: maps a game to its tabular Q-learning state. - -Referenced from babysnake/pyproject.toml's [tool.retro-gamer] section, and -used directly by train_babysnake.py via GameEnvironment. -""" - -ACTIONS = ["KEY_RIGHT", "KEY_DOWN", "KEY_LEFT", "KEY_UP"] - - -def get_state(game): - """Return BabySnake's state as a hashable (agent_x, agent_y, food_x, food_y) tuple.""" - s = game.state - return (int(s['agent_x']), int(s['agent_y']), int(s['food_x']), int(s['food_y'])) diff --git a/q_learning.py b/q_learning.py index 3e599bb..7f54d29 100644 --- a/q_learning.py +++ b/q_learning.py @@ -1,79 +1,99 @@ -"""Q-learning agent. - -This file contains starter code for implementing a Q-learning agent. -You need to fill in two functions: - - choose_action: select an action using an epsilon-greedy policy - - update_q: update the Q-table using the Bellman equation - -Both functions are written here with no reference to BabySnake or any -particular game — they work on any (state, action) Q-table, given a list of -possible actions. train_babysnake.py is the module that wires these functions -up to BabySnake specifically. - -Before training, get your implementation passing the tests in -test_q_learning.py: - - python test_q_learning.py - -Errors caught here are much easier to track down than errors discovered -during training. Once the tests pass, run: - - python train_babysnake.py -""" - import random +from itertools import count +from retro_gamer import GameEnvironment, GameMetadata - -def choose_action(q_table, state, actions, epsilon): - """Choose an action using an epsilon-greedy policy. - - With probability `epsilon`, return a random action from `actions`. - Otherwise, return the action with the highest Q-value in `q_table` - for the given `state`. If a (state, action) pair has not been seen - before, treat its Q-value as 0. +class QLearning: + """ + Trains a policy to play a game using Q-learning. + The Q-table maintains an estimate of the quality of every action at every state, + and the estimates are improved through training. Arguments: - q_table (dict): Maps (state, action) -> Q-value. - state: The current state. - actions (list): The actions available to choose from. - epsilon (float): Exploration rate, between 0.0 and 1.0. - - Returns: - One of the values in `actions`. - - Hint: random.random() returns a float in [0.0, 1.0). - random.choice(actions) returns a random action. - q_table.get(key, default) is handy for missing entries. + 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. """ - raise NotImplementedError("Fill in choose_action") + def __init__( + self, episodes=1000, alpha=0.1, gamma=0.95, epsilon=1.0, + epsilon_decay=0.995, epsilon_min=0.05, max_steps_per_episode=500 + ): + self.episodes = episodes + self.alpha = alpha + self.gamma = gamma + self.epsilon = epsilon + self.epsilon_decay = epsilon_decay + self.epsilon_min = epsilon_min + self.max_steps_per_episode = max_steps_per_episode -def update_q(q_table, state, action, reward, next_state, actions, alpha, gamma): - """Update one entry of the Q-table using the Bellman equation. + def train(self, env, actions): + """Trains the policy by updating estimates of the quality of state/actions + in the Q-table. - The update rule is: + Arguments: + env (retro_gamer.GameEnvironment): a game environment, which allows + stepping through a game one action at a time. + """ + self.Q = {} + for episode in range(self.episodes): + state = env.reset() + turn = 0 + total_reward = 0 + for turn in range(self.max_steps_per_episode): + action = self.choose_action(state, actions) + next_state, reward, done = env.step(action) + self.update_q(state, action, reward, next_state, actions) + state = next_state + total_reward += reward + self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay) + self.report_progress(episode, total_reward) + return self.Q - Q(s, a) <- Q(s, a) + alpha * (r + gamma * max_a' Q(s', a') - Q(s, a)) + def choose_action(self, state, actions): + """Choose an action using an epsilon-greedy policy. + With probability `self.epsilon`, return a random action from `actions`. + Otherwise, return the action with the highest Q-value in `self.Q` + for the given `state`. If a (state, action) pair has not been seen + before, treat its Q-value as 0. + """ + if random.random() < self.epsilon: + return random.choice(actions) + else: + action_qualities = [(self.Q.get((state, a), 0), a) for a in actions] + best_q, best_action = sorted(action_qualities, reverse=True)[0] + return best_action - where: - s, a — the state we were in and the action we took - r — the reward we received - s' — the state we ended up in - max_a' ... — the best possible Q-value from the new state + def update_q(self, state, action, reward, next_state, next_actions): + """Update an entry in self.Q. The update rule is: - Arguments: - q_table (dict): Maps (state, action) -> Q-value (modified in place). - state: The state before the action. - action: The action taken. - reward (float): The reward received. - next_state: The state after the action. - actions (list): The actions available from `next_state`. - alpha (float): Learning rate (how much to update). - gamma (float): Discount factor (how much to value future rewards). + Q(s, a) <- Q(s, a) + alpha * (r + gamma * max_a' Q(s', a') - Q(s, a)) - Returns: - None — modifies q_table in place. + where: + s, a the state we were in and the action we took + r the reward we received + s' the next state + max_a' Q(s', a') the best possible Q-value of actions from the next state + alpha the learning rate (use self.alpha) + gamma the discount factor (use self.gamma) + """ + future_qs = [(self.Q.get((next_state, a), 0), a) for a in next_actions] + best_future_q, best_future_action = sorted(future_qs, reverse=True)[0] + q = self.Q.get((state, action), 0) + self.Q[(state, action)] = q + self.alpha * (reward + self.gamma * best_future_q - q) + + def report_progress(self, episode, total_reward): + if (episode + 1) % 100 == 0: + print( + f"Episode {episode + 1:5d} " + f"reward={total_reward:6.1f} " + f"epsilon={self.epsilon:.3f} " + f"q_entries={len(self.Q)}" + ) - Hint: Q-values for unseen (state, action) pairs start at 0. - """ - raise NotImplementedError("Fill in update_q") diff --git a/q_learning_solution.py b/q_learning_solution.py deleted file mode 100644 index f510906..0000000 --- a/q_learning_solution.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Solution for q_learning.py — remove before publishing to students.""" - -import random - - -def choose_action(q_table, state, actions, epsilon): - if random.random() < epsilon: - return random.choice(actions) - q_values = [q_table.get((state, a), 0.0) for a in actions] - return actions[q_values.index(max(q_values))] - - -def update_q(q_table, state, action, reward, next_state, actions, alpha, gamma): - old_q = q_table.get((state, action), 0.0) - next_q_values = [q_table.get((next_state, a), 0.0) for a in actions] - best_next_q = max(next_q_values) - new_q = old_q + alpha * (reward + gamma * best_next_q - old_q) - q_table[(state, action)] = new_q diff --git a/runs/snake/checkpoints/ep_2300.pt b/runs/snake/checkpoints/ep_2300.pt deleted file mode 100644 index 4929f6b..0000000 Binary files a/runs/snake/checkpoints/ep_2300.pt and /dev/null differ diff --git a/train_babysnake.py b/train_babysnake.py index a7b99ca..ed50783 100644 --- a/train_babysnake.py +++ b/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)