Files
lab_reinforcement_learning/q_learning.py
Chris Proctor e8a24ae7be Split q_learning.py into algorithm, environment glue, and a training script
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.
2026-06-24 07:50:28 -04:00

80 lines
2.8 KiB
Python

"""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
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.
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.
"""
raise NotImplementedError("Fill in choose_action")
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:
Q(s, a) <- Q(s, a) + alpha * (r + gamma * max_a' Q(s', a') - Q(s, a))
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
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).
Returns:
None — modifies q_table in place.
Hint: Q-values for unseen (state, action) pairs start at 0.
"""
raise NotImplementedError("Fill in update_q")