Refactoring lab

This commit is contained in:
Chris Proctor
2026-06-26 13:27:11 -04:00
parent e752bb848b
commit a9385f8296
32 changed files with 644 additions and 686 deletions

0
games/__init__.py Normal file
View File

126
games/babysnake/__init__.py Normal file
View File

@@ -0,0 +1,126 @@
"""BabySnake: a 4×4 grid game where an agent collects food.
State: (agent_x, agent_y, food_x, food_y) — four integers.
The agent starts with 50 energy. Each step costs 1 energy.
Collecting food restores 30 energy and adds 1 to the score.
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 = 8
START_ENERGY = 50
FOOD_ENERGY = 30
class Forager:
"""The player agent. Controlled with arrow keys."""
name = "Forager"
character = '@'
color = "green_on_black"
position = (0, 0)
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)
def __init__(self):
self._direction = self.RIGHT
def handle_keystroke(self, keystroke, game):
if keystroke.name == "KEY_RIGHT":
self._direction = self.RIGHT
elif keystroke.name == "KEY_UP":
self._direction = self.UP
elif keystroke.name == "KEY_LEFT":
self._direction = self.LEFT
elif keystroke.name == "KEY_DOWN":
self._direction = self.DOWN
def play_turn(self, game):
x, y = self.position
dx, dy = self._direction
new_pos = (x + dx, y + dy)
if game.on_board(new_pos):
self.position = new_pos
game.state['energy'] -= 1
game.state['reward'] -= 0.01
food = game.get_agent_by_name("Food")
if self.position == food.position:
food.relocate(game)
game.state['energy'] += FOOD_ENERGY
game.state['score'] += 1
game.state['reward'] += 1.0
ax, ay = self.position
fx, fy = game.get_agent_by_name("Food").position
game.state['agent_x'] = ax
game.state['agent_y'] = ay
game.state['food_x'] = fx
game.state['food_y'] = fy
if game.state['energy'] <= 0:
game.end()
class Food:
"""The food item. Respawns at a random empty position when collected."""
name = "Food"
character = '*'
color = "yellow_on_black"
position = (0, 0)
def relocate(self, game):
bw, bh = game.board_size
forager = game.get_agent_by_name("Forager")
while True:
pos = (randint(0, bw - 1), randint(0, bh - 1))
if pos != forager.position:
self.position = pos
return
def create_game():
"""Return a fresh BabySnake game."""
forager = Forager()
food = Food()
bw = bh = BOARD_SIZE
game = Game(
[forager, food],
{
'score': 0,
'reward': 0.0,
'energy': START_ENERGY,
'agent_x': 0,
'agent_y': 0,
'food_x': 0,
'food_y': 0,
},
board_size=(bw, bh),
framerate=6,
show_state=False,
)
forager.position = (randint(0, bw - 1), randint(0, bh - 1))
food.relocate(game)
ax, ay = forager.position
fx, fy = food.position
game.state['agent_x'] = ax
game.state['agent_y'] = ay
game.state['food_x'] = fx
game.state['food_y'] = fy
return game
if __name__ == '__main__':
create_game().play()

View File

@@ -0,0 +1,3 @@
from games.babysnake import create_game
create_game().play()

View File

@@ -0,0 +1,11 @@
[project]
name = "babysnake"
version = "0.1.0"
description = "BabySnake: a tiny grid game for teaching Q-learning"
requires-python = ">=3.11"
dependencies = ["retro-games>=2.5.0"]
[tool.retro-gamer]
actions = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]
reward = "reward"
observation_function = "games.babysnake:get_state"

41
games/babysnake/train.py Normal file
View File

@@ -0,0 +1,41 @@
"""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)

161
games/frogger/__init__.py Normal file
View File

@@ -0,0 +1,161 @@
"""Frogger: guide a frog across busy traffic lanes to reach the far side.
The frog starts at the bottom row. Cars move across the lanes between the
start and goal. Each lane has cars moving at different speeds in alternating
directions. The frog earns +10 for each row advanced, +50 for reaching the
top row, and -10 for being hit by a car or falling off the edge. Episodes end
when the frog reaches the top, gets hit, or energy runs out.
"""
from random import randint
from retro.game import Game
BOARD_WIDTH = 20
BOARD_HEIGHT = 12
NUM_LANES = BOARD_HEIGHT - 2
START_ENERGY = 200
class Frog:
name = "Frog"
character = 'O'
color = "green_on_black"
position = (0, 0)
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)
def __init__(self):
self._direction = self.UP
def handle_keystroke(self, keystroke, game):
if keystroke.name == "KEY_UP":
self._direction = self.UP
elif keystroke.name == "KEY_DOWN":
self._direction = self.DOWN
elif keystroke.name == "KEY_LEFT":
self._direction = self.LEFT
elif keystroke.name == "KEY_RIGHT":
self._direction = self.RIGHT
def play_turn(self, game):
bw, bh = game.board_size
x, y = self.position
dx, dy = self._direction
nx, ny = x + dx, y + dy
if not (0 <= nx < bw):
game.state['reward'] -= 10
game.state['energy'] -= 50
self._reset(game)
return
if not (0 <= ny < bh):
if ny < 0:
game.state['score'] += 50
game.state['reward'] += 50
else:
game.state['reward'] -= 5
self._reset(game)
return
prev_y = y
self.position = (nx, ny)
game.state['energy'] -= 1
game.state['reward'] -= 0.01
if ny < prev_y:
advancement = prev_y - ny
game.state['score'] += advancement * 10
game.state['reward'] += advancement * 5
for agent in game.agents:
if hasattr(agent, '_is_car') and agent.position == self.position:
game.state['reward'] -= 10
game.state['energy'] -= 50
self._reset(game)
return
bw, bh = game.board_size
fx, fy = self.position
game.state['frog_x'] = fx / bw
game.state['frog_y'] = fy / bh
if game.state['energy'] <= 0:
game.end()
def _reset(self, game):
bw, bh = game.board_size
self.position = (bw // 2, bh - 1)
self._direction = self.UP
if game.state['energy'] <= 0:
game.end()
class Car:
_is_car = True
character = 'X'
color = "red_on_black"
def __init__(self, lane, speed, direction, start_x, board_width):
self.name = f"Car {lane}_{start_x}"
self._lane = lane
self._speed = speed
self._direction = direction
self._board_width = board_width
self._step = 0
self.position = (start_x, lane)
def play_turn(self, game):
self._step += 1
if self._step < self._speed:
return
self._step = 0
x, y = self.position
x = (x + self._direction) % self._board_width
self.position = (x, y)
frog = game.get_agent_by_name("Frog")
if frog.position == self.position:
frog._reset(game)
game.state['reward'] -= 10
game.state['energy'] -= 50
def create_game():
bw, bh = BOARD_WIDTH, BOARD_HEIGHT
frog = Frog()
frog.position = (bw // 2, bh - 1)
agents = [frog]
car_id = 0
for lane_idx, row in enumerate(range(1, bh - 1)):
direction = 1 if lane_idx % 2 == 0 else -1
speed = 2 + (lane_idx % 3)
num_cars = 2 + (lane_idx % 3)
spacing = bw // num_cars
for i in range(num_cars):
start_x = (i * spacing + lane_idx * 3) % bw
agents.append(Car(row, speed, direction, start_x, bw))
car_id += 1
game = Game(
agents,
{
'score': 0,
'reward': 0.0,
'energy': START_ENERGY,
'frog_x': (bw // 2) / bw,
'frog_y': (bh - 1) / bh,
},
board_size=(bw, bh),
framerate=8,
show_state=['score', 'energy'],
)
return game
if __name__ == '__main__':
create_game().play()

View File

@@ -0,0 +1,3 @@
from games.frogger import create_game
create_game().play()

View File

@@ -0,0 +1,11 @@
[project]
name = "frogger"
version = "0.1.0"
description = "Frogger: guide a frog across traffic lanes to train an RL agent"
requires-python = ">=3.11"
dependencies = ["retro-games>=2.5.0"]
[tool.retro-gamer]
actions = ["KEY_UP", "KEY_DOWN", "KEY_LEFT", "KEY_RIGHT"]
reward = "reward"
character_set = ["O", "X"]

36
games/snake/__init__.py Normal file
View File

@@ -0,0 +1,36 @@
from random import randint, choice
from .snake import SnakeHead
from .apple import Apple
from retro.game import Game
def create_game():
"""Return a fresh, initialized Snake game."""
head = SnakeHead()
apple = Apple()
game = Game(
[head, apple],
{'score': 0, 'reward': 0, 'energy': 150, 'apple_dx': 0.0, 'apple_dy': 0.0},
board_size=(32, 16),
framerate=12,
)
bw, bh = game.board_size
head.position = (randint(1, bw - 2), randint(1, bh - 2))
direction, character = choice([
(SnakeHead.RIGHT, '>'),
(SnakeHead.UP, '^'),
(SnakeHead.LEFT, '<'),
(SnakeHead.DOWN, 'v'),
])
head.direction = direction
head.character = character
apple.relocate(game)
hx, hy = head.position
ax, ay = apple.position
game.state['apple_dx'] = (ax - hx) / bw
game.state['apple_dy'] = (ay - hy) / bh
return game
if __name__ == '__main__':
create_game().play()

3
games/snake/__main__.py Normal file
View File

@@ -0,0 +1,3 @@
from games.snake import create_game
create_game().play()

49
games/snake/apple.py Normal file
View File

@@ -0,0 +1,49 @@
from random import randint
class Apple:
"""An agent representing the Apple.
Note how Apple doesn't have ``play_turn`` or
``handle_keystroke`` methods: the Apple doesn't need to do
anything in this game. It just sits there waiting to get
eaten.
Attributes:
name: "Apple"
character: '@'
color: "red_on_black" (`Here's documentation on how colors
work <https://blessed.readthedocs.io/en/latest/colors.html>`_
position: (0, 0). The Apple will choose a random position
as soon as the game starts, but it needs an initial
position to be assigned.
"""
name = "Apple"
character = '@'
color = "red_on_black"
position = (0, 0)
def relocate(self, game):
"""Sets position to a random empty position. This method is
called whenever the snake's head touches the apple.
Arguments:
game (Game): The current game.
"""
self.position = self.random_empty_position(game)
def random_empty_position(self, game):
"""Returns a randomly-selected empty position. Uses a very
simple algorithm: Get the game's board size, choose a
random x-value between 0 and the board width, and choose
a random y-value between 0 and the board height. Now use
the game to check whether any Agents are occupying this
position. If so, keep randomly choosing a new position
until the position is empty.
"""
bw, bh = game.board_size
occupied_positions = game.get_agents_by_position()
while True:
position = (randint(0, bw-1), randint(0, bh-1))
if position not in occupied_positions:
return position

View File

@@ -0,0 +1,18 @@
import numpy as np
from retro.views.headless import HeadlessView
from retro_gamer.observation import egocentric_board, encode_board, encode_state
CHARACTER_SET = ["@", "*", ">", "<", "^", "v"]
RADIUS = 8
def egocentric_observation(game):
"""17×17 window centered on the snake's head, plus apple_dx and apple_dy."""
view = HeadlessView()
view.on_game_start(game)
view.render(game)
head = game.get_agent_by_name("Snake head")
cropped = egocentric_board(view.board_characters, head.position, RADIUS)
board_vec = encode_board(cropped, CHARACTER_SET).flatten()
extras = encode_state(game.state, ["apple_dx", "apple_dy"])
return np.concatenate([board_vec, extras])

View File

@@ -0,0 +1,10 @@
[project]
name = "snake"
version = "0.1.0"
description = "Snake: a classic game for training RL agents"
requires-python = ">=3.11"
dependencies = ["retro-games>=2.5.0"]
[tool.retro-gamer]
actions = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]
reward = "reward"

160
games/snake/snake.py Normal file
View File

@@ -0,0 +1,160 @@
class SnakeHead:
"""An Agent representing the snake's head. When the game starts, you control
the snake head using the arrow keys. The SnakeHead always has a direction, and
will keep moving in that direction every turn. When you press an arrow key,
you change the SnakeHead's direction.
Attributes:
name: "Snake head"
position: (0,0)
character: ``'v'`` Depending on the snake head's direction, its character
changes to ``'<'``, ``'^'``, ``'>'``, or ``'v'``.
next_segment: Initially ``None``, this is a reference to a SnakeBodySegment.
growing: When set to True, the snake will grow a new segment on its next move.
"""
RIGHT = (1, 0)
UP = (0, -1)
LEFT = (-1, 0)
DOWN = (0, 1)
name = "Snake head"
position = (0, 0)
direction = DOWN
character = 'v'
next_segment = None
growing = False
def play_turn(self, game):
"""On each turn, the snake head uses its position and direction to figure out
its next position. If the snake head is able to move there (it's on the board and
not occuppied by part of the snake's body), it moves.
Then, if the snake head is on the Apple, the Apple moves to a new random position
and ``growing`` is set to True.
Now we need to deal with two situations. First, if ``next_segment`` is not None, there is
a SnakeBodySegment attached to the head. We need the body to follow the head,
so we call ``self.next_segment.move``, passing the head's old position
(this will be the body's new position), a reference to the game, and a value for
``growing``. If the snake needs to grow, we need to pass this information along
the body until it reaches the tail--this is where the next segment will be attached.
If there is no ``next_segment`` but ``self.growing`` is True, it's time to add
a body! We set ``self.next_segment`` to a new SnakeBodySegment, set its
position to the head's old position, and add it to the game. We also add 1 to the
game's score.
"""
x, y = self.position
dx, dy = self.direction
next_pos = (x + dx, y + dy)
if self.can_move(next_pos, game):
apple = game.get_agent_by_name("Apple")
ax, ay = apple.position
old_dist = abs(x - ax) + abs(y - ay)
new_dist = abs(next_pos[0] - ax) + abs(next_pos[1] - ay)
game.state['reward'] += old_dist - new_dist # +1 toward apple, -1 away
game.state['energy'] -= 1
self.position = next_pos
if self.is_on_apple(self.position, game):
apple.relocate(game)
self.growing = True
game.state['score'] += 50
game.state['reward'] += 50
game.state['energy'] = 150
if self.next_segment:
self.next_segment.move((x, y), game, growing=self.growing)
elif self.growing:
self.next_segment = SnakeBodySegment(1, (x, y))
game.add_agent(self.next_segment)
self.growing = False
bw, bh = game.board_size
hx, hy = self.position
ax, ay = apple.position
game.state['apple_dx'] = (ax - hx) / bw
game.state['apple_dy'] = (ay - hy) / bh
if game.state['energy'] <= 0:
game.state['reward'] -= 10
game.end()
else:
game.state['reward'] -= 10
game.end()
def handle_keystroke(self, keystroke, game):
"""Checks whether one of the arrow keys has been pressed.
If so, sets the SnakeHead's direction and character.
"""
x, y = self.position
if keystroke.name == "KEY_RIGHT":
self.direction = self.RIGHT
self.character = '>'
elif keystroke.name == "KEY_UP":
self.direction = self.UP
self.character = '^'
elif keystroke.name == "KEY_LEFT":
self.direction = self.LEFT
self.character = '<'
elif keystroke.name == "KEY_DOWN":
self.direction = self.DOWN
self.character = 'v'
def can_move(self, position, game):
on_board = game.on_board(position)
empty = game.is_empty(position)
on_apple = self.is_on_apple(position, game)
return on_board and (empty or on_apple)
def is_on_apple(self, position, game):
apple = game.get_agent_by_name("Apple")
return apple.position == position
class SnakeBodySegment:
"""Finally, we need an Agent for the snake's body segments.
SnakeBodySegment doesn't have ``play_turn`` or ``handle_keystroke`` methods because
it never does anything on its own. It only moves when the SnakeHead, or the previous
segment, tells it to move.
Arguments:
segment_id (int): Keeps track of how far back this segment is from the head.
This is used to give the segment a unique name, and also to keep track
of how many points the player earns for eating the next apple.
position (int, int): The initial position.
Attributes:
character: '*'
next_segment: Initially ``None``, this is a reference to a SnakeBodySegment
when this segment is not the last one in the snake's body.
"""
character = '*'
next_segment = None
def __init__(self, segment_id, position):
self.segment_id = segment_id
self.name = f"Snake body segment {segment_id}"
self.position = position
def move(self, new_position, game, growing=False):
"""When SnakeHead moves, it sets off a chain reaction, moving all its
body segments. Whenever the head or a body segment has another segment
(``next_segment``), it calls that segment's ``move`` method.
This method updates the SnakeBodySegment's position. Then, if
``self.next_segment`` is not None, calls that segment's ``move`` method.
If there is no next segment and ``growing`` is True, then we set
``self.next_segment`` to a new SnakeBodySegment in this segment's old
position, and update the game's score.
Arguments:
new_position (int, int): The new position.
game (Game): A reference to the current game.
growing (bool): (Default False) When True, the snake needs to
add a new segment.
"""
old_position = self.position
self.position = new_position
if self.next_segment:
self.next_segment.move(old_position, game, growing=growing)
elif growing:
self.next_segment = SnakeBodySegment(self.segment_id + 1, old_position)
game.add_agent(self.next_segment)

View File

@@ -0,0 +1,92 @@
from random import randint, choice
from games.snake.apple import Apple
from games.snake.snake import SnakeBodySegment
from retro.game import Game
class SnakeHead:
RIGHT = (1, 0)
UP = (0, -1)
LEFT = (-1, 0)
DOWN = (0, 1)
name = "Snake head"
position = (0, 0)
direction = DOWN
character = 'v'
next_segment = None
growing = False
def play_turn(self, game):
x, y = self.position
dx, dy = self.direction
next_pos = (x + dx, y + dy)
if self.can_move(next_pos, game):
apple = game.get_agent_by_name("Apple")
ax, ay = apple.position
old_dist = abs(x - ax) + abs(y - ay)
new_dist = abs(next_pos[0] - ax) + abs(next_pos[1] - ay)
game.state['reward'] += old_dist - new_dist
game.state['energy'] -= 1
self.position = next_pos
if self.is_on_apple(self.position, game):
apple.relocate(game)
self.growing = True
game.state['score'] += 50
game.state['reward'] += 50
game.state['energy'] = 150
if self.next_segment:
self.next_segment.move((x, y), game, growing=self.growing)
elif self.growing:
self.next_segment = SnakeBodySegment(1, (x, y))
game.add_agent(self.next_segment)
self.growing = False
if game.state['energy'] <= 0:
game.state['reward'] -= 10
game.end()
else:
game.state['reward'] -= 10
game.end()
def handle_keystroke(self, keystroke, game):
if keystroke.name == "KEY_RIGHT":
self.direction = self.RIGHT
self.character = '>'
elif keystroke.name == "KEY_UP":
self.direction = self.UP
self.character = '^'
elif keystroke.name == "KEY_LEFT":
self.direction = self.LEFT
self.character = '<'
elif keystroke.name == "KEY_DOWN":
self.direction = self.DOWN
self.character = 'v'
def can_move(self, position, game):
return game.on_board(position) and (game.is_empty(position) or self.is_on_apple(position, game))
def is_on_apple(self, position, game):
return game.get_agent_by_name("Apple").position == position
def create_game():
head = SnakeHead()
apple = Apple()
game = Game(
[head, apple],
{'score': 0, 'reward': 0, 'energy': 150},
board_size=(32, 16),
framerate=12,
show_state=['score'],
)
bw, bh = game.board_size
head.position = (randint(1, bw - 2), randint(1, bh - 2))
direction, character = choice([
(SnakeHead.RIGHT, '>'),
(SnakeHead.UP, '^'),
(SnakeHead.LEFT, '<'),
(SnakeHead.DOWN, 'v'),
])
head.direction = direction
head.character = character
apple.relocate(game)
return game

View File

@@ -0,0 +1,10 @@
[project]
name = "snake_v1"
version = "0.1.0"
description = "Snake (v1): full board, no extra features"
requires-python = ">=3.11"
dependencies = ["retro-games>=2.5.0"]
[tool.retro-gamer]
actions = ["KEY_RIGHT", "KEY_UP", "KEY_LEFT", "KEY_DOWN"]
reward = "reward"