Snake AI

A small reinforcement learning project that trains and compares two Deep Q-Network (DQN) agents to play Snake in a headless 20×20 grid environment.

This project includes:

• a custom Snake game environment
• a pure DQN agent with experience replay and a target network
• a heuristic-enhanced DQN agent that avoids immediate death and prefers safer moves
• a side-by-side visual demo for comparing both agents
• training scripts to save model checkpoints to the models/ folder

Project overview

The environment is implemented in game.py and models the classic Snake game as a reinforcement learning task.

• Grid size: 20 × 20
• State vector: 11 features describing danger, movement direction, and food direction
• Actions: 0 = straight, 1 = turn right, 2 = turn left
• Reward design:
  • eating food: +10
  • moving closer to food: +0.1
  • moving farther from food: -0.1
  • crash / starvation: large negative penalty

The agent logic is split between:

• model.py: neural network architecture for the DQN
• agent.py: replay buffer, pure DQN agent, and heuristic DQN agent
• train.py: training loop and model saving
• demo.py: pygame visualization for playing the agents side by side

Repository structure

snake/
├── agent.py          # DQN agents and replay buffer
├── demo.py           # Pygame visual comparison demo
├── game.py           # Snake environment logic
├── model.py          # Neural network model definition
├── models/           # Trained model checkpoints
├── requirements.txt  # Python dependencies
├── train.py          # Training script
└── README.md         # Project documentation

Features

Pure DQN

The PureDQNAgent uses a basic DQN architecture with:

• epsilon-greedy exploration
• replay memory
• target network updates
• Bellman-style TD learning

This agent learns from rewards with no domain-specific guidance.

Heuristic DQN

The HeuristicDQNAgent extends the same model with two safety heuristics:

1. Death masking: it avoids moves that would instantly crash if another safe move exists.
2. Flood-fill bias: when multiple safe actions are available, it prefers the move with more open space.

This improves training quality and gives the demo a stronger-looking agent.

Setup

Create a virtual environment and install dependencies:

python -m venv .venv
source .venv/bin/activate       # macOS / Linux
# or .venv\Scripts\activate     # Windows

pip install -r requirements.txt

Training

Train both models with the default settings:

python train.py

Train only the pure DQN agent:

python train.py --pure-only

Train only the heuristic DQN agent:

python train.py --heuristic-only

Customize the number of episodes:

python train.py --episodes 800

The models are saved into:

• models/pure_dqn.pth
• models/heuristic_dqn.pth

Demo

Run the visual comparison demo:

python demo.py

The demo shows two games side by side:

• left panel: Pure DQN
• right panel: Heuristic DQN

Controls:

• Up / Down arrow: increase or decrease game speed
• Space: pause / resume
• Esc: quit

Notes

• The demo will exit early with a message if the saved model files are missing, so it is best to train once before launching the visualizer.
• Training is designed to run in a headless environment and is intentionally lightweight enough for CPU usage.
• The project uses pygame for visualization, torch for deep learning, and numpy for state handling.

Typical workflow

python train.py
python demo.py

This gives you a trained Snake AI baseline and a visual comparison between raw DQN learning and heuristic-assisted DQN learning.