Step 10 — ML Extension 🤖

The Big Idea: Learning by Trial and Error

Instead of a human pressing keys, an AI agent chooses actions by looking up a table of learned values (the Q-table). Over thousands of games, it learns which actions lead to reward and which lead to death.

State Representation

The Q-learner needs to know the current situation (state). We can't use raw coordinates — the grid is too large. Instead we use 11 boolean features that capture what matters.

def get_state(snake, food, direction):
    """
    Return an 11-element tuple of booleans representing the game state.
    This tuple becomes the KEY in our Q-table dictionary.
    """
    hx, hy = snake.head()
    # 3 danger sensors: is there immediate danger straight/left/right?
    dir_r = (direction == "RIGHT")
    dir_l = (direction == "LEFT")
    dir_u = (direction == "UP")
    dir_d = (direction == "DOWN")

    # Check danger in relative directions
    def is_wall_or_body(x, y):
        """True if (x, y) is out of bounds or snake body."""
        if x < 0 or x >= GRID_COLS or y < 0 or y >= GRID_ROWS:
            return True
        return (x, y) in snake.body[1:]

    state = (
        # Danger straight, right, left (relative to current direction)
        is_wall_or_body(hx + (dir_r - dir_l), hy + (dir_d - dir_u)),  # danger ahead
        is_wall_or_body(hx + (dir_d - dir_u), hy - (dir_r - dir_l)),  # danger right
        is_wall_or_body(hx - (dir_d - dir_u), hy + (dir_r - dir_l)),  # danger left

        # Current movement direction (one-hot encoded)
        dir_r, dir_l, dir_u, dir_d,

        # Food direction relative to head
        food.pos[0] < hx,   # food is to the left
        food.pos[0] > hx,   # food is to the right
        food.pos[1] < hy,   # food is above
        food.pos[1] > hy,   # food is below
    )
    return state   # a tuple of 11 True/False values

🧠 State Vector — Interactive Visualizer

The Q-learner sees the game as 11 True/False features. Click the arrow buttons to move the snake head. Watch which features flip as danger and food directions change — this tuple becomes the Q-table key.

The Q-Table

The Q-table maps (state, action) pairs to a numeric value — how good is it to take this action in this situation? Higher Q-value = better action.

Example Q-table (simplified)

The agent picks the action with the highest Q-value for the current state. In Python, this table is just a defaultdict:

from collections import defaultdict
import random

# Q-table: maps (state, action) → float
# defaultdict(float) returns 0.0 for any new key
q_table = defaultdict(float)

ACTIONS = [0, 1, 2]   # 0=straight, 1=turn-left, 2=turn-right
EPSILON = 0.1         # exploration rate (10% random, 90% best known)

def choose_action(state):
    """ε-greedy: explore randomly 10% of the time, exploit best 90%."""
    if random.random() < EPSILON:
        return random.choice(ACTIONS)   # explore!
    q_vals = [q_table[(state, a)] for a in ACTIONS]
    return ACTIONS[q_vals.index(max(q_vals))]   # exploit!

The Q-Learning Update Rule

After every action, we update the Q-value using the Bellman equation. New estimate = immediate reward + discounted future reward.

ALPHA = 0.1    # learning rate: how fast we update Q values
GAMMA = 0.9    # discount factor: how much we value future rewards

def update_q(state, action, reward, next_state):
    """
    Q-Learning update rule (Bellman equation):
    Q(s,a) ← Q(s,a) + α × [r + γ × max(Q(s',a')) − Q(s,a)]
    """
    old_q = q_table[(state, action)]
    # Best possible future reward from next_state
    best_future = max(q_table[(next_state, a)] for a in ACTIONS)
    # Bellman target
    target = reward + GAMMA * best_future
    # Update Q-value (weighted average with learning rate)
    q_table[(state, action)] = old_q + ALPHA * (target - old_q)

# Reward signal:
# +10 = ate food   (reward good behaviour)
# -10 = died       (punish bad behaviour)
# -0.1 = each step (small penalty encourages efficiency)

The Training Loop

Instead of one human playing, we run thousands of games in a loop. Each game updates the Q-table. After enough games, the agent learns to play well.

import pygame
from collections import defaultdict

# (All class definitions from Step 8 go here)

q_table = defaultdict(float)
ALPHA = 0.1; GAMMA = 0.9; EPSILON = 0.1

# Map relative action (0,1,2) to absolute direction
TURNS = {
    "RIGHT": ["RIGHT","DOWN","UP"],
    "LEFT":  ["LEFT","UP","DOWN"],
    "UP":    ["UP","RIGHT","LEFT"],
    "DOWN":  ["DOWN","LEFT","RIGHT"],
}

def train(n_episodes=10000, visual=False):
    """Train the Q-learner over n_episodes games."""
    pygame.init()
    screen = pygame.display.set_mode(
        (GRID_COLS*CELL_SIZE, GRID_ROWS*CELL_SIZE)) if visual else None
    clock = pygame.time.Clock()
    scores = []

    for episode in range(n_episodes):
        game = Game()
        state = get_state(game.snake, game.food, game.snake.direction)
        steps = 0

        while game.running and steps < 1000:   # cap steps to avoid loops
            action = choose_action(state)
            new_dir = TURNS[game.snake.direction][action]
            game.snake.direction = new_dir

            old_score = game.score
            game.update()
            steps += 1

            if game.score > old_score:
                reward = 10   # ate food
            elif not game.running:
                reward = -10  # died
            else:
                reward = -0.1  # still alive, small penalty

            next_state = get_state(game.snake, game.food,
                                   game.snake.direction)
            update_q(state, action, reward, next_state)
            state = next_state

            if visual:
                game.draw(); clock.tick(60)

        scores.append(game.score)
        if episode % 1000 == 0:
            avg = sum(scores[-100:])/min(100, len(scores))
            print(f"Episode {episode:5d} | Avg score: {avg:.2f}"
                  f" | Q-table size: {len(q_table)}")

    pygame.quit()
    return scores

if __name__ == "__main__":
    train(n_episodes=5000, visual=True)

What Happens During Training?

Early in training the snake dies quickly (Q-table is all zeros, actions are random). After a few thousand games, patterns emerge.

The Big Picture — Where This Leads

Try It — Exercises

🎉 You've completed the Snake Game Tutorial!

You went from print("Hello, World!") to building and training an AI agent — in 9 steps. That's remarkable.

← Back to the Tutorial Index

⚠️ Common Mistakes — Spot These Before You Start

These are the errors beginners make most often in Step 10. Read them now so you can recognise them in your own code.

LEARNING_RATE = 1.0   # replaces Q-value completely each update

LEARNING_RATE = 0.1   # blend 10% new info with 90% existing knowledge

EPSILON = 0.3   # fixed: agent always explores 30% of the time

epsilon = max(0.01, epsilon * 0.995)  # decay: explore less as agent learns

# Trying to store the state IN the action variable

state = get_state(snake, food)   # observe the world
action = choose_action(state)    # decide what to do
q_table[state][action] += ...    # update the value of that choice

✅ Quick Check — Are You Ready?

Three questions — not graded. They help you spot gaps before the activities.