mwc/lab_reinforcement_learning
Template
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Babysnake training works

Browse Source
This commit is contained in:
Chris Proctor
2026-06-25 13:00:03 -04:00
parent aeb610d04b
commit 048cb1c02b
7 changed files with 105 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
babysnake/__init__.py
View File

@@ -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

2
babysnake/pyproject.toml
View File

@@ -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"

13
babysnake_env.py
View File

@@ -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']))

152
q_learning.py
View File

@@ -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
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
def update_q(self, state, action, reward, next_state, next_actions):
"""Update an entry in self.Q. 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.
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)
"""
raise NotImplementedError("Fill in update_q")
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)}"
)

18
q_learning_solution.py
View File

@@ -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

BIN
runs/snake/checkpoints/ep_2300.pt
View File

Binary file not shown.

83
train_babysnake.py
View File

@@ -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)