diff --git a/q_learning.py b/q_learning.py index 8ab3011..3e599bb 100644 --- a/q_learning.py +++ b/q_learning.py @@ -1,103 +1,54 @@ -"""Q-learning agent for BabySnake. +"""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 -Run this file to train the agent: - python q_learning.py +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. -After training, run this to watch it play: - python -c "from q_learning import watch; watch()" +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 -import babysnake -from retro.input import ProgrammaticInput -from retro.views.headless import HeadlessView - -# The four actions the agent can take. -ACTIONS = ["KEY_RIGHT", "KEY_DOWN", "KEY_LEFT", "KEY_UP"] -# --------------------------------------------------------------------------- -# Environment wrapper -# --------------------------------------------------------------------------- - -class BabySnakeEnv: - """A simple wrapper that lets us step through BabySnake programmatically. - - Usage: - env = BabySnakeEnv() - state = env.reset() # start a new episode - next_state, reward, done = env.step("KEY_RIGHT") - """ - - def reset(self): - """Start a new episode. Returns the initial state tuple.""" - self._inp = ProgrammaticInput() - self.game = babysnake.create_game() - self.game.input_source = self._inp - self.game.view = HeadlessView() - self.game.start() - self._prev_reward = 0.0 - return self._get_state() - - def step(self, action): - """Take one action. Returns (next_state, reward, done). - - Arguments: - action (str): One of ACTIONS, or None for no-op. - - Returns: - next_state (tuple): The state after the action. - reward (float): The reward received this step. - done (bool): True if the episode has ended. - """ - self._inp.press(action) - self.game.step() - next_state = self._get_state() - reward = self.game.state['reward'] - self._prev_reward - self._prev_reward = self.game.state['reward'] - done = not self.game.playing - return next_state, reward, done - - def _get_state(self): - """Return the current state as a tuple of four integers.""" - s = self.game.state - return (int(s['agent_x']), int(s['agent_y']), - int(s['food_x']), int(s['food_y'])) - - -# --------------------------------------------------------------------------- -# Q-learning functions — fill these in! -# --------------------------------------------------------------------------- - -def choose_action(q_table, state, epsilon): +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. + 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. Arguments: q_table (dict): Maps (state, action) -> Q-value. - state (tuple): The current state, e.g. (1, 2, 3, 0). + state: The current state. + actions (list): The actions available to choose from. epsilon (float): Exploration rate, between 0.0 and 1.0. Returns: - str: One action from ACTIONS. + One of the values in `actions`. Hint: random.random() returns a float in [0.0, 1.0). - random.choice(ACTIONS) returns a random action. + random.choice(actions) returns a random action. q_table.get(key, default) is handy for missing entries. """ raise NotImplementedError("Fill in choose_action") -def update_q(q_table, state, action, reward, next_state, alpha, gamma): +def update_q(q_table, state, action, reward, next_state, actions, alpha, gamma): """Update one entry of the Q-table using the Bellman equation. The update rule is: @@ -112,10 +63,11 @@ def update_q(q_table, state, action, reward, next_state, alpha, gamma): Arguments: q_table (dict): Maps (state, action) -> Q-value (modified in place). - state (tuple): The state before the action. - action (str): The action taken. + state: The state before the action. + action: The action taken. reward (float): The reward received. - next_state (tuple): The state after the action. + 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). @@ -125,94 +77,3 @@ def update_q(q_table, state, action, reward, next_state, alpha, gamma): Hint: Q-values for unseen (state, action) pairs start at 0. """ raise NotImplementedError("Fill in update_q") - - -# --------------------------------------------------------------------------- -# Training loop -# --------------------------------------------------------------------------- - -def train( - episodes=1000, - alpha=0.1, - gamma=0.95, - epsilon=1.0, - epsilon_decay=0.995, - epsilon_min=0.05, -): - """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. - - Returns: - dict: The trained Q-table. - """ - q_table = {} - env = BabySnakeEnv() - - for episode in range(episodes): - state = env.reset() - total_reward = 0.0 - - while env.game.playing: - action = choose_action(q_table, state, epsilon) - next_state, reward, done = env.step(action) - update_q(q_table, state, action, reward, next_state, 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. - """ - import babysnake - from retro.input import ProgrammaticInput - - 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): - s = game.state - state = (int(s['agent_x']), int(s['agent_y']), - int(s['food_x']), int(s['food_y'])) - q_values = [q_table.get((state, a), 0.0) for a in ACTIONS] - best = ACTIONS[q_values.index(max(q_values))] - _inp.press(best) - 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) diff --git a/q_learning_solution.py b/q_learning_solution.py index a70700c..f510906 100644 --- a/q_learning_solution.py +++ b/q_learning_solution.py @@ -1,113 +1,18 @@ """Solution for q_learning.py — remove before publishing to students.""" import random -import babysnake -from retro.input import ProgrammaticInput -from retro.views.headless import HeadlessView - -ACTIONS = ["KEY_RIGHT", "KEY_DOWN", "KEY_LEFT", "KEY_UP"] -class BabySnakeEnv: - def reset(self): - self._inp = ProgrammaticInput() - self.game = babysnake.create_game() - self.game.input_source = self._inp - self.game.view = HeadlessView() - self.game.start() - self._prev_reward = 0.0 - return self._get_state() - - def step(self, action): - self._inp.press(action) - self.game.step() - next_state = self._get_state() - reward = self.game.state['reward'] - self._prev_reward - self._prev_reward = self.game.state['reward'] - done = not self.game.playing - return next_state, reward, done - - def _get_state(self): - s = self.game.state - return (int(s['agent_x']), int(s['agent_y']), - int(s['food_x']), int(s['food_y'])) - - -def choose_action(q_table, state, epsilon): +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))] + 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, alpha, gamma): +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] + 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 - - -def train( - episodes=1000, - alpha=0.1, - gamma=0.95, - epsilon=1.0, - epsilon_decay=0.995, - epsilon_min=0.05, -): - q_table = {} - env = BabySnakeEnv() - - for episode in range(episodes): - state = env.reset() - total_reward = 0.0 - - while env.game.playing: - action = choose_action(q_table, state, epsilon) - next_state, reward, done = env.step(action) - update_q(q_table, state, action, reward, next_state, 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): - if q_table is None: - print("Training first...") - q_table = train() - - _inp = ProgrammaticInput() - - class PolicyInput: - def collect(self): - s = game.state - state = (int(s['agent_x']), int(s['agent_y']), - int(s['food_x']), int(s['food_y'])) - q_values = [q_table.get((state, a), 0.0) for a in ACTIONS] - best = ACTIONS[q_values.index(max(q_values))] - _inp.press(best) - 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) diff --git a/questions.md b/questions.md new file mode 100644 index 0000000..a40aabd --- /dev/null +++ b/questions.md @@ -0,0 +1,51 @@ +# Questions + +## BabySnake + +## Checkpoint 1: Before training + +1. How do you decide where to move in BabySnake? Explain how to choose moves + in enough detail that someone else could follow your instructions. + +2. How many distinct states are there for BabySnake? If we assume that all four + arrow keys are valid actions in every state, how many rows would the full Q-table contain? + +3. The discount factor γ (gamma) can range from 0 to 1. What would be the effect of setting + γ to 0? What about 1? + +4. The learning rate α (alpha) can also range from 0 to 1. What would be the effect of setting + α to 0? What about 1? + +5. Calculate the new Q-value for the situation described. Explain your answer. + +6. Implement `choose_action` and `update_q` in `q_learning.py`, then run + + ``` + python test_q_learning.py + ``` + + Get every test passing before moving on — errors are much easier to track + down here than during training. + + +--- + +## Checkpoint 2: After training + +Train your Q-learning agent to consistently score 3 or more food items per +episode, then watch it play: + +``` +python train_babysnake.py +``` + +**At what episode did the agent start reliably finding food?** + + +**Print `q_table` after training. Can you read the policy?** For a state you +pick, does the highest Q-value point toward the food? + + +**How does the trained agent's behavior compare to the reasoning you wrote +down in Checkpoint 1?** + diff --git a/test_q_learning.py b/test_q_learning.py new file mode 100644 index 0000000..76c931c --- /dev/null +++ b/test_q_learning.py @@ -0,0 +1,62 @@ +# test_q_learning.py +# ------------ +# Defines tests for `q_learning`. Run this program with `python test_q_learning.py`. +# You don't need to edit this file. +# +# Get every test here passing before you run train_babysnake.py — errors are +# much easier to spot here than during training. + +from unittest import TestCase, main +import random + +from q_learning import choose_action, update_q + +ACTIONS = ["UP", "DOWN", "LEFT", "RIGHT"] + + +class TestChooseAction(TestCase): + def test_greedy_picks_highest_q_value(self): + q_table = {("s", "UP"): 1.0, ("s", "DOWN"): 5.0, ("s", "LEFT"): 2.0, ("s", "RIGHT"): 0.0} + self.assertEqual(choose_action(q_table, "s", ACTIONS, epsilon=0.0), "DOWN") + + def test_unseen_state_defaults_to_zero_and_picks_first_action(self): + self.assertEqual(choose_action({}, "new_state", ACTIONS, epsilon=0.0), ACTIONS[0]) + + def test_fully_random_explores_more_than_one_action(self): + random.seed(0) + q_table = {("s", "UP"): 100.0} # UP is clearly the best action + results = {choose_action(q_table, "s", ACTIONS, epsilon=1.0) for _ in range(50)} + self.assertGreater(len(results), 1) + + def test_always_returns_a_valid_action(self): + for _ in range(20): + result = choose_action({}, "s", ACTIONS, epsilon=0.5) + self.assertIn(result, ACTIONS) + + +class TestUpdateQ(TestCase): + def test_basic_bellman_update(self): + q_table = {("s", "UP"): 0.0} + update_q(q_table, "s", "UP", reward=1.0, next_state="t", actions=ACTIONS, alpha=0.5, gamma=0.9) + # old_q=0, best_next_q=0 (unseen) -> target=1.0, new_q = 0 + 0.5*(1.0-0) = 0.5 + self.assertAlmostEqual(q_table[("s", "UP")], 0.5) + + def test_uses_best_next_q_value(self): + q_table = {("s", "UP"): 0.0, ("t", "UP"): 2.0, ("t", "DOWN"): 5.0} + update_q(q_table, "s", "UP", reward=0.0, next_state="t", actions=ACTIONS, alpha=1.0, gamma=1.0) + # target = 0 + 1.0*5.0 = 5.0; alpha=1 fully replaces the old value + self.assertAlmostEqual(q_table[("s", "UP")], 5.0) + + def test_alpha_zero_means_no_change(self): + q_table = {("s", "UP"): 3.0} + update_q(q_table, "s", "UP", reward=10.0, next_state="t", actions=ACTIONS, alpha=0.0, gamma=0.9) + self.assertAlmostEqual(q_table[("s", "UP")], 3.0) + + def test_only_updates_the_given_state_action_pair(self): + q_table = {} + update_q(q_table, "s", "UP", reward=1.0, next_state="t", actions=ACTIONS, alpha=0.5, gamma=0.9) + self.assertEqual(set(q_table.keys()), {("s", "UP")}) + + +if __name__ == '__main__': + main() diff --git a/train_babysnake.py b/train_babysnake.py new file mode 100644 index 0000000..a7b99ca --- /dev/null +++ b/train_babysnake.py @@ -0,0 +1,104 @@ +"""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)