It Learned Snake By Dying 85,000 Times

Most "AI plays a game" demos you've seen aren't learning anything. They're a decision tree wearing a neural-network costume: if wall_ahead: turn, if food_is_right: go_right, dressed up and called intelligence. This one isn't that. It's a DQN (Deep Q-Network) — a small neural net that starts out knowing nothing about walls, food, or its own tail, and learns all three from nothing but a +10 when it eats and a -10 when it dies.

There is no avoid_wall() function anywhere in this codebase. Nobody wrote a chase_food() heuristic. The agent found both behaviors on its own, the same way you'd learn a new game if someone only ever told you "good" or "bad" after each move and never explained why.

Full code: github.com/rash-aad/snake. Let's open it up.

The Five Pieces

Everything downstream depends on one decision made early: keep the game logic completely dumb and completely fast. game.py doesn't know pygame exists. It doesn't know a neural network exists. It's a few hundred lines of grid math — move a head, check a collision, hand back a score — and because it does nothing else, it can run tens of thousands of games per second during training instead of the handful per second a rendered game would allow.

Everything else is a layer bolted on top of that.

• game.py — pure Snake rules. No rendering, no RL, no opinions.
• env.py — the translator. Turns the raw board into a 17-number state vector the network can read, and turns the network's 3 possible outputs into an actual direction.
• agent.py — the DQN itself: the network, the replay buffer, the target network, epsilon-greedy exploration.
• train.py — the loop that actually runs the whole thing thousands of times and saves progress.
• play.py — the only file that imports pygame. You can play it yourself, or run python play.py --watch to load the trained weights and watch the agent play, with a live panel showing what it's "thinking."

Nothing here is exotic. It's the standard DQN recipe from the 2015 Atari paper, scaled down to a problem small enough to actually reason about by hand.

What The Network Actually Sees

The single most important design decision in this whole project isn't the network — it's env.py's _get_state(), because a DQN can only ever be as good as the information you hand it. Give it a bad state representation and no amount of training fixes that.

The first version of this only checked the one cell directly ahead in three directions: straight, left, right. That's enough to not walk into a wall on the very next step, but it's blind to anything more than one square away, including the snake's own tail curling back toward its head from two or three cells out. The fix was 8 distance sensors, one per compass direction, each reporting how much open space lies between the head and the nearest wall or body segment, normalized to 0..1:

The two blocked rays make the point. W is the obvious case — a naive "check the next cell" sensor would catch that one too. NW is the one that actually justified the upgrade: the tail curling back in from two cells away, diagonally, invisible to anything that only looks one step ahead. That's genuine body awareness, not just wall awareness.

_COMPASS = [(0,-1), (1,-1), (1,0), (1,1), (0,1), (-1,1), (-1,0), (-1,-1)]

for dx, dy in _COMPASS:
    dist = 0
    x, y = head.x, head.y
    while True:
        x += dx; y += dy; dist += 1
        if out_of_bounds(x, y): break
        if Point(x, y) in body: break
    sensors.append(dist / grid_size)

The other 9 numbers in the 17-number state round it out: 4 for current direction (one-hot: left/right/up/down), 4 for which side the food is on relative to the head, and 1 for the snake's own length, normalized to 0..1 — so the network can tell "I'm 4 cells long, plenty of room" from "I'm 180 cells long, this board is basically all tail now" and play more cautiously as it grows.

The Network Itself

17 numbers go in. Three numbers come out: one score per possible move (go straight, turn right, turn left — all relative to current heading, so "reverse into your own neck" isn't even in the action space to begin with). Those three numbers are Q-values — the network's own estimate of "how much total future reward do I expect if I take this action from this state."

The whole network is two lines of PyTorch — Linear(17, 256) -> ReLU -> Linear(256, 3) — and the interesting part isn't the architecture, it's how the targets it's trained against get built:

q_new = reward if done else reward + GAMMA * max(target_model(next_state))

That's the Bellman equation, and it's the entire mechanism by which "I ate food 40 moves from now" turns into a signal that reaches the decision made right now: reward propagates backward through the GAMMA * max(...) term, one training step at a time, until a move that merely sets up a future opportunity starts scoring higher than a move that doesn't, even though neither one collects any reward on the spot.

Two stability tricks make that actually work in practice instead of diverging:

• A replay buffer (100,000 transitions, sampled in batches of 1,000) training on a random batch of old moves instead of only the move that just happened, so the network isn't learning from a stream of highly-correlated back-to-back frames.
• A target network — a frozen copy of the model, used only to compute the max(...) in the formula above, synced to the live weights every 25 games. Without it, the network is chasing a target generated by the same weights it's currently updating — a moving goalpost that shifts under it every single step.

Exploration follows a simple linear decay: epsilon = max(0, 80 - n_games), compared against a random(0, 200) roll. At game 0 that's a 40% chance of a random move; by game 80 it's exactly 0%, and every move after that is argmax(Q) — pure exploitation of whatever the network has learned so far.

Watching It Think

play.py --watch doesn't just render the board — it prints the network's reasoning in plain English, ranked by Q-value, right below the game:

Under the hood it's honest, not decorative: the three candidate moves get ranked by the network's raw Q-values, and each one gets a plain-English description grounded in the actual board state, not the network's internals:

if game.is_collision(new_head):
    return "would crash immediately"
return "moves toward food" if new_dist < cur_dist else "moves away from food"

The chosen move (argmax(Q)) gets an arrow next to it. Watching this panel for a few minutes is a better intuition-builder for what a Q-value actually is than any equation: it's not a plan, it's a running vote, recomputed from scratch every single frame, over three futures the network has never explicitly been told exist.

What Broke, And What Fixed It

A checkpoint bug that silently ate progress. For a long stretch, the model was only saved to disk on a new high score. That sounds reasonable until you hit a plateau — thousands of games with no new record — during which killing and resuming the training process meant "resuming" from a snapshot that was, unknowingly, far older than the process that just died. The fix was blunt on purpose: save a checkpoint every 50 games regardless of whether it's a record, so a kill/resume never loses more than a few dozen games again.

Reading noise as signal, twice. Early on, a mean score computed from just 5 games right after a restart got mistaken for a converged average, and a real architectural conclusion ("reward shaping is hurting performance") got built on top of that misread before anyone noticed the sample size was the problem, not the reward function. Nothing about the code changed to fix this. the fix was discipline: a mean only means something once it's averaged over hundreds of games, not five.

A real ceiling, confirmed instead of assumed. After roughly 85,000 games, the model plateaued hard: record stuck at 59, mean stuck around 20, for tens of thousands of games regardless of what else got tweaked. It would've been easy to keep nudging hyperparameters and call it fixed the moment the mean twitched upward. Instead the plateau got confirmed across multiple checkpoints before it got accepted as real: this specific architecture — 17 inputs, one 256-wide hidden layer — has a genuine capacity limit, and no additional training time on its own was ever going to move it further.

Watching It Play

The trained model reliably scores in the 15-25 range per game, record 59, reacting sensibly to both food position and its own body — not a scripted response, a policy that emerged from a few hundred thousand rounds of "was that good or bad." It won't clear the entire 20×20 board; a guaranteed full clear is a fundamentally different, non-learned problem (think Hamiltonian cycle, not reward shaping). But within the rules it was actually trained on, it plays.

The Whole Idea, In One Line

Give a network a compact, honest description of the world (8 rays, a direction, a food bearing, a length) and a single scalar telling it whether the last move helped or hurt, and it will, slowly and with a few very real bugs along the way, build its own understanding of walls and tails and food without ever being told those concepts exist. That's not a trick specific to Snake — it's the same loop (state -> action -> reward -> update) underneath every DQN, on every game, at every scale.

Full source, train.log, and the saved model.pth are on GitHub: github.com/rash-aad/snake. Clone it, delete model.pth, and watch it relearn all of this from a completely random network, in real time.

Eco signing off