Snake RL

Overview

The environment is a headless-friendly pygame implementation (snake_env.py): a 640×480 grid with a discrete action space of three relative moves — straight, turn right, turn left — and rewards of +10 for food, −10 for dying, and 0 otherwise. Moving into the tail cell is legal (it vacates on the same tick), and an idle cap of 100 × snake length steps since the last food truncates the episode instead of looping forever.

The agent (model.py) is a dueling Double DQN: a 28 → 256 → 256 ReLU body, then separate value and advantage heads summed dueling-style. It trains against a frozen target network (hard update every 1000 steps) with 3-step returns, Huber loss, and an ε-greedy schedule decaying from 1.0 to 0.005 over 150k steps.

The state (state.py) encodes 28 egocentric features: a full obstacle ray per move direction, flood-fill free space per candidate move (what lets the agent avoid sealing itself into a dead end), food and tail position in the snake's own frame, and board occupancy. The v1 state was an 11-feature encoding that could not see its own body shape and plateaued at a mean of ~18.

Results

Measured on the same 32×24 board: greedy evaluation of runs/v2/best.pth averages 102 per game (max 143 in 20 eval games, 173 in a 50-game run) — roughly a 5× jump over the v1 mean of ~18. The checkpoint is saved whenever the greedy eval mean improves, so it is protected against a late-training dip.

The trained weights are also exported to src/lib/rl/snake-dqn-weights.json so the same agent can play in the browser — a TypeScript port of the env, state encoding, and dueling forward pass (no training). Watch it play live.

Run it yourself

uv run train.py   # train a new agent headless (~1500 steps/s)
uv run play.py    # watch the saved best.pth play
uv run plot.py    # render checkpoints/log.csv -> progress.png

Dependencies: torch, pygame, numpy — pinned in pyproject.toml (Python ≥ 3.12, managed with uv). Training is headless by default; add --render to watch it learn.

Notes & next steps

• v1 → v2: the original 11-feature state could only see the three cells touching the head, so a long snake had no way to perceive that it was sealing itself in — the state was not Markov enough to support scores past ~30 no matter how long you trained. The 28-feature encoding (rays, flood-fill free space, egocentric food/tail geometry) is what unlocks the ~5× jump.
• Algorithm fixes: the v1 update never detached its bootstrap target, so gradients flowed into it and the net regressed onto itself. v2 uses a frozen target network (hard update every 1000 steps), Double DQN action selection, a dueling head, 3-step returns, and Huber loss with grad clipping.
• Next step: the current encoding is hand-designed. The real jump is a convolutional net over a stacked grid representation, which removes feature engineering entirely — or a Hamiltonian-cycle solver for near-perfect play.