Reinforcement Learning

In this lab, you will train a computer to play games—without telling it the rules. Instead, you will set up a system where the computer tries things, observes what happens, and gradually learns which actions lead to better outcomes. This is called reinforcement learning.

You have probably already seen how machine learning models learn to classify data from labeled examples. Reinforcement learning is different: there is no labeled dataset and no right answer to compare against. The agent learns entirely from the consequences of its own actions.

Training BabySnake

💻 Run the BabySnake game:

$ python -m games.babysnake

BabySnake is a simple game on an 8×8 grid. You control the @ character and try to collect the * food. Use the arrow keys to move. The game starts with 50 energy, loses 1 energy per step, and ends when energy runs out.

Play several rounds and pay attention to how you decide where to move.

Policies

Every game can be represented as a set of states, and each state has one or more possible actions. When you select an action, you end up in a new state, and you get a reward. A policy is a function which decides what to do in every possible state--you wrote an informal policy in your answer to question 1. The goal of reinforcement learning is to discover a good policy (one which ends up with the most reward) automatically, from experience rather than from a rule you write by hand.

One way to represent a policy is Q-table, a table of states, actions, and the estimated quality (Q) of each. The quality of a (state, action) pair is an estimate of all the future rewards you would get if you took that action from that state. If you have an accurate Q-table for a game, then at every state, you could just compare the quality of the possible actions and choose the one with the highest quality.

For BabySnake, the state is a tuple of four integers: (agent_x, agent_y, food_x, food_y). Actions are the arrow keys. The reward is +1 for eating food, -0.01 otherwise. Consider this state on a simplified 4×4 board:

@ * . .
. . . .
. . . .
. . . .

Here's our Q-table. There are two actions available for this state. Moving right has a higher quality than moving down, so the agent should move right. This example is tediously obvious, but for a more complex game, a Q-table can encode all the skill of an expert player.

Q-learning

So how do we get a good Q-table for a game? Have an agent play the game a lot and learn from experience. Start by setting every Q-value to 0. Then, every time the agent is at state $s$, plays an action $a$, receives a reward $r$, and reaches a new state $s'$, update the Q-value of $(s, a)$ according to this rule:

$$ Q(s, a) \leftarrow Q(s, a) + \alpha \left( r + \gamma \max_{a'} Q(s', a') - Q(s, a) \right) $$

Let's break it down, from the inside out:

• $Q(s, a)$ is the quality of action $a$ at state $s$. The $\leftarrow$ means we are going to assign a new value to $Q(s, a)$, like = in Python.
• After the agent at $s$ plays action $a$, receives reward $r$, and arrives at $s'$, we can make a better estimate of $Q(s, a)$: $r + \gamma \max_{a'} Q(s', a')$, or the reward we received plus the expected quality of the best possible action $a'$ we could take from $s'$. (We can find $\max_{a'} Q(s', a')$ by looking up all the rows of the Q-table with $s'$ and choosing the highest quality.) We apply a discount factor $\gamma$ (the Greek letter gamma) to future rewards, so that rewards now are worth more than rewards in the future.
• If we subtract the old Q-value $Q(s, a)$ from the new Q-value $r + \gamma \max_{a'} Q(s', a')$, we get the temporal difference error, or the amount by which the old Q-value was wrong.
• Multiply the temporal difference error by the learning rate $\alpha$, and add this to the old Q-value. Assign this as the new Q-value.

Exploration vs. exploitation

During training, the agent faces a dilemma at every turn: exploit what it already knows (choose the action with the highest Q-value), or explore by trying something new (maybe there is a better action it hasn't discovered yet)?

We handle this with an epsilon-greedy policy: with probability ε (epsilon), take a random action (explore); with probability 1 − ε, take the best known action (exploit). At the start of training, ε is high (e.g., 1.0: always random). As training progresses, ε decays toward a small floor (e.g., 0.05), so the agent explores widely early on and exploits what it has learned later.

Implement Q-learning

💻 Open q_learning/__init__.py. The QLearning class is mostly written for you. You need to implement two methods:

choose_action(self, state, actions)

Choose one of the actions and return it. With probability self.epsilon, the agent should explore (return a random action from actions). Otherwise, the agent should exploit (choose the best action from actions by comparing their Q-values).

• random.random() returns a float between 0 and 1.
• random.choice(actions) returns a random action.

update_q(self, state, action, reward, next_state, next_actions)

Once an action has been taken from a state and we have the reward and the next state, we can learn from the experience. Apply the update rule to update self.Q[(state, action)].

• Use self.Q.get((state, action), 0) to look up a Q-value, with 0 as the default in case (state, action) is not yet a key in the dict.
• Use self.alpha and self.gamma.

💻 Verify your implementation before training:

$ python -m q_learning.tests

All 8 tests should pass. Then train the agent and watch it play:

$ python -m games.babysnake.train

Press Enter or Escape to quit the game viewer.

Training Snake

Now let's move on to the real snake game.

💻 First, play the game by running python -m games.snake.

╔════════════════════════════════╗
║                                ║
║                                ║
║                                ║
║                                ║
║                                ║
║                                ║
║                                ║
║                                ║
║                ***             ║
║                *               ║
║                *               ║
║                *               ║
║            <****               ║
║    @                           ║
║                                ║
╚════════════════════════════════╝

You will notice a few upgrades: the board is much larger, the snake grows every time you eat an apple, and the snake dies when it collides with a wall or its own body.

This presents some challenges for reinforcement learning. It's no longer enough to represent the game state as the location of the snake head and the apple, we need to keep track of the snake's whole body. Considering the larger board size, there are trillions of possible states. A Q-table representing all these states would not fit in your computer's memory, and it would take centuries of training to encounter even a fraction of the possible states.

What we need instead is something that generalizes: given a state it has never exactly seen before, it should still produce a reasonable estimate, based on similar states it has seen.

Deep Q-Networks

A Q-network is a neural network that plays the same role as a Q-table: it takes an observation of the game and outputs one Q-value per action. Training adjusts its weights so its outputs get closer to the optimal Q-value: the reward for the current action plus the discounted value of the best subsequent action. The gap between the network's prediction and this target is the loss; training aims to minimize the loss, just like in the neural classification lab. This training process is called Deep Q-Learning (DQN).

We won't get into all of the details here; instead, we'll use a high-level library, retro-gamer, which uses DQN to train agents to play any games written using retro-games. In this section we'll explore a pretrained policy for Snake--you could easily re-do the training yourself (retro-gamer clean runs/snake-v1 and then retro-gamer train runs/snake-v1), it would just take a while.

The goal is to develop a high-level understanding of how the game state representation, the structure of the neural network, and the training process affect a successful training process. If you continue learning about machine learning in the future, you will spend most of your time working at this high level using libraries such as PyTorch, which let you build neural networks and then handle all the math for you.

Attempt 1: The full board

We started by setting up a new training by running:

$ retro-gamer init games.snake runs/snake-v1

This created the runs/snake-v1 directory, which contains the following:

runs/snake-v1
├── checkpoints
│   ├── ep_0100.pt
│   ├── ...
│   └── ep_5000.pt
├── config.toml
└── training.log

The most important file here is config.toml, which specifies parameters for the training. retro-gamer reads the game's metadata in its pyproject.toml file, and uses this to fill in important details such as the board size, the allowed actions, how to determine the reward at each state, etc. See the documentation for all the details.

By default, the game's state is just what's on the board. The game's character_set specifies all the characters that might appear. For Snake, this is ["@", "*", ">", "<", "^", "v"] (the apple, the snake body, and the four possible orientation of the snake's head). Each cell is represented by a list where every value is 0 except a single 1. A cell containing the apple would be encoded as [1, 0, 0, 0, 0, 0], and a cell containing a snake body segment (*) would be encoded as [0, 1, 0, 0, 0, 0]. This is called "one-hot" encoding.

Snake's full board is 32 cells wide and 16 cells high, and each cell requires six numers when hot-encoded. Therefore, Snake's state is represented as $32 * 16 * 6 = 3072$ numbers, each 0 or 1. Our DQN represents Q(state, action), so this state, along with an encoded action, will be the input to a neural network exactly like those we studied in the Classification:Neural lab; the number of hidden layers and their sizes are specified in config.toml. The output of the network is an estimate of the quality of the action at the state, and the training process will update the values of the hidden layers so that the network produces better and better estimates.

It's a good idea to reward your agent for good behavior. In this version of Snake, we reward the snake with 1 point whenever it moves toward the apple, -1 point when it moves away from the apple, and 50 points when it eats an apple.

training.log contains a record of the training process. For each batch of 100 episodes, the log shows the average reward, the average number of steps the snake survived, and other hyperparameters which change during training.

[ep_0100]  ep=0001-0100  avg_reward=-11.9  avg_steps= 50  epsilon=0.905  avg_loss=0.6  time=0m08s  total=0m08s
[ep_0500]  ep=0401-0500  avg_reward= +4.9  avg_steps= 91  epsilon=0.606  avg_loss=2.2  time=0m17s  total=1m00s
[ep_0900]  ep=0801-0900  avg_reward=+14.7  avg_steps=100  epsilon=0.406  avg_loss=2.6  time=0m34s  total=2m53s
[ep_5000]  ep=4901-5000  avg_reward=+27.6  avg_steps=100  epsilon=0.050  avg_loss=1.4  time=0m28s  total=22m43s

💻 Plot the training parameters with:

$ retro-gamer plot runs/snake-v1

After 5,000 episodes, reward was still noisy and plateauing around 20–30, which is fine but not great. The agent learned to survive—steps per episode rose from 50 to around 100—but finding apples remained unreliable.

Each checkpoint is a saved copy of the neural network; you can load any of them up to see the agent play using that policy. (We have not included all the checkpoints in the training runs, because checkpoints take up a fair amount of space.)

💻 Watch the fully-trained v1 agent play (press Enter or Escape to quit):

$ retro-gamer play runs/snake-v1 --checkpoint ep_5000

Attempt 2: An arrow to the apple

What if the agent didn't have to discover "direction to the apple" on its own—what if we just told it? We added two numbers to the observation (the game updates these values in its state dict every turn):

• apple_dx = (apple_x − head_x) / board_width
• apple_dy = (apple_y − head_y) / board_height

Together, these numbers are a vector pointing from the snake's head toward the apple. In snake-v2, we add these numbers to the state representation in config.toml:

observe_state = [
    "apple_dx",
    "apple_dy",
]

Once again, plot the training results and watch the agent play.

Within 500 episodes, reward had reached 43, more than Attempt 1 managed after thousands of episodes. But look at average steps: as reward rose, episodes got shorter. By episode 1,000, the average episode was only 32 steps long. With no information about walls or the snake's growing body, the agent learned to charge toward the apple—and crashed often.

💻 Watch the v2 agent—notice what it does well and where it goes wrong:

$ retro-gamer play runs/snake-v2 --checkpoint ep_3000

If you watch the snake crash a few times, you may notice that it often crashes immediately after it eats an apple, then turns back onto itself as it aims for the next apple.

Attempt 3: A convolutional network

Attempt 2 succeeded by giving the agent a shortcut—it never had to discover the apple's direction from the board. But it paid a price: with no information about walls or its own body, the snake charges straight at the apple and crashes into its own body. Episodes stayed short (around 32 steps), and the snake never learned to navigate around itself.

What if we gave it the whole board again, but used a better architecture? A convolutional neural network (CNN) processes a 2-D grid by learning small spatial filters—patterns like "wall one step ahead" or "body nearby"—that apply everywhere on the board. Instead of treating the 32×16 board as a flat list of 3,072 numbers, it scans the board with overlapping 3×3 windows and builds up spatial feature maps. This is how image recognition models work, and a game board is just a small image.

In snake-v3, we add the board back and switch on the CNN:

[preprocessing]
spatial = true
observe_state = ["apple_dx", "apple_dy"]

The trainer logs the architecture it built:

CNN: Conv2d(6→32, k=3, pad=1) → ReLU → Conv2d(32→64, k=3, pad=1) → ReLU
MLP head: 32770 → 128 → 64 → 5

💻 Plot the results:

$ retro-gamer plot runs/snake-v3

Reward rose quickly to around +30 by episode 1,900—then stalled. After epsilon reached its minimum around episode 3,000, reward stopped improving and ended at +21.2 at episode 5,000. Training took 1h 22m—twenty times longer than Attempt 2, for a result significantly worse.

This is a regression, and regressions happen. The CNN can see the body and walls, but it has to discover useful spatial patterns entirely on its own from a large, mostly empty board. The network is far bigger, training is far slower, and the extra information turns out to be more of a burden than a help.

💻 Watch the agent at its peak (episode 1,900) and compare its behavior to v2:

$ retro-gamer play runs/snake-v3 --checkpoint ep_1900

Attempt 4: An egocentric CNN

Attempt 3 didn't fail because CNNs are wrong for Snake—it failed because the framing was wrong. Feeding the entire 32×16 board to a CNN means the network has to learn to locate the snake inside a mostly-empty grid before it can even start reasoning about what's around it. Most of the 3,072 input values are blank space with no useful signal.

We decided to try an egocentric observation: instead of showing the whole board, crop a 17×17 window centered on the snake's head. The head is always at the center of the window; what varies is what's around it—walls, body segments, the apple. Now the CNN's local filters are looking at locally meaningful patterns, the same way they work in image recognition.

We implement this as a custom observation function:

# games/snake/observation.py
def egocentric_cnn_observation(game):
    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=8)
    board_vec = encode_board(cropped, CHARACTER_SET).transpose(2, 0, 1).flatten()
    extras = encode_state(game.state, ["apple_dx", "apple_dy"])
    return np.concatenate([board_vec, extras])

The config points to it and sets board = false so the default board encoding doesn't become part of the input.

[metadata]
observation_function = "games.snake.observation:egocentric_cnn_observation"

[preprocessing]
spatial = true
board = false

The trainer builds the same CNN architecture, now over a 17×17 board:

CNN: Conv2d(6→32, k=3, pad=1) → ReLU → Conv2d(32→64, k=3, pad=1) → ReLU
MLP head: 18498 → 128 → 64 → 5

💻 Plot the results:

$ retro-gamer plot runs/snake-v4

By episode 400, reward had already reached +26—a level that took Attempt 3 until episode 1,400. By episode 3,000, it had passed +100, more than Attempt 2 ever reached. At 5,000 episodes (our initial setting for training_episodes), reward was still climbing, so we updated training_episodes to 10,000 and kept training. Finally, the CNN converged to around +270, peaking at +287 — in 53 minutes total.

The only thing that changed was the observation. We used the same CNN and the same hyperparameters, but we cropped the board view to the area around the snake. This may have made it easier for the CNN to learn patterns like "lots of snake body near the center is dangerous."

💻 Watch the v4 agent:

$ retro-gamer play runs/snake-v4 --checkpoint ep_10000

Attempt 5: A narrow local view

If a 17×17 egocentric window works well with a CNN, does the window even need to be that large? And do we need a CNN at all? A 7×7 window—three cells in every direction from the head—captures what's immediately relevant: the wall ahead, the body segment to the left, the apple if it's nearby. With only 49 cells, the board part of the observation is just 294 numbers, small enough for a plain MLP.

# games/snake/observation.py
def narrow_egocentric_observation(game):
    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=3)
    board_vec = encode_board(cropped, CHARACTER_SET).flatten()
    extras = encode_state(game.state, ["apple_dx", "apple_dy"])
    return np.concatenate([board_vec, extras])

Since we don't want the standard board encoding to be part of the input, we set board = false.

[metadata]
observation_function = "games.snake.observation:narrow_egocentric_observation"

[preprocessing]
board = false

💻 Plot the results:

$ retro-gamer plot runs/snake-v5

Reward climbed steeply—+48 by episode 500, +100 by episode 900—and kept going, stabilizing around +250 after episode 3,000. Final reward: +258.5 in just 12 minutes. A network with 296 inputs, trained in under a quarter of an hour, slightly outperformed the 17×17 CNN.

💻 Watch the v5 agent and notice how it handles its growing body:

$ retro-gamer play runs/snake-v5 --checkpoint ep_5000

Training Frogger

Now it's your turn. games/frogger/ is a Frogger-style game: a frog starts at the bottom of a board and tries to cross busy traffic lanes to reach the top. Cars move back and forth across the lanes; the frog earns points for each successful crossing of the road within 60 seconds.

💻 Play Frogger to get a feel for it:

$ python -m games.frogger

Use the arrow keys to move the frog. Press Enter or Escape to quit.

Setting up a training run

💻 Create your first training run:

$ retro-gamer init games/frogger runs/frogger

Initialized training run at runs/frogger/config.toml
  game        : games/frogger
  board_size  : 20×12
  actions     : ['KEY_UP', 'KEY_DOWN', 'KEY_LEFT', 'KEY_RIGHT']
  reward      : reward
  characters  : ['O', 'X']
  architecture: MLP (non-spatial)

💻 Start training. You can stop training at any time with Ctrl-C and resume later by running the same train command again.

$ retro-gamer train runs/frogger

💻 At any point, watch your agent play or check its progress:

$ retro-gamer play runs/frogger
$ retro-gamer plot runs/frogger

Run three attempts, starting with the default configuration. If you want to compare runs side by side, create separate directories (e.g. runs/frogger-v2) the same way you initialized the first one. Here are some things you might experiment with:

• epsilon_decay — how quickly should the agent commit to what it has learned?
• learning_rate — too high risks instability; too low learns very slowly.
• gamma — how much should the agent care about future rewards vs. immediate ones?
• hidden_sizes in [model] — a larger network can represent more complex policies at the cost of training speed.
• observe_state under [preprocessing] — does telling the agent its own position or other game-state values help it learn faster?

When you change hidden_sizes or anything under [preprocessing], run retro-gamer clean runs/frogger first—these changes make existing checkpoints incompatible, and retro-gamer train will refuse to resume.

Extension

If you've built a game with the retro-games framework, you can train a DQN agent to play it. The retro-gamer walkthrough in its documentation is a good place to start.

Closing discussion: imagining new applications

This is the last lab in the AI unit -- estimation, classification (features, then neural), and now reinforcement learning. Before moving on, take stock of the toolkit: predicting a number from data, sorting things into categories by hand-picked or learned features, and learning a policy from trial and error and a reward signal.