Reinforcement learning with snake in pygame

RL building blocks

Agent

The agent is an entity that can enforce actions on the environment and observe changes in it. According to the actions it takes, it is rewarded or punished in the process of training. It's learning itself through trial and error and receives rewards as unique feedback. In our case the agent is replacing the human player in the process of controlling the snake.

Configurable parameters:

• MAX_MEM - maximal number of sets of states stored for the training
• BATCH_SIZE - size of single batch to train on
• LR - learning rate

Perception

The agent is able to sniff food in 4 main directions - left,up,right,down. Sniffing is based on position of the fruit obtained from the game environment and projected as boolean list of directions.

Environment

The environment consists the table, the apple and the snake. They provide the information about current state and rewards for the actions taken by the agent. In our case the environment is represented by following classes: GameAI, SnakeAI, Fruit and Table. The input for environment is agent's action and the outputs are observation and reward for the step or set of steps.

Configurable parameters:

• SCREEN_WIDTH - table width
• SCREEN_HEIGHT - table height
• TICK - how fast should the snake move
• SCALE - how wide the snake should be

Action

The action is taken by the agent and is affecting the environment, more specifically the snake's next move. In our case action is represented as a list of 0's and 1's according to the following schema: [straight, right, left].

Example:

• [1,0,0] means go straight wrt to the snake's head direction (choosing purple arrow)
• [0,1,0] means go straight wrt to the snake's head direction (choosing orange arrow)
• [0,0,1] means go straight wrt to the snake's head direction (choosing cyan arrow)

Reward

The reward is assigned based on the performance of the agent in the environment. The agent gets positive reward for getting the apple and negative for hitting walls/himself or exceeding the time intended for exploration.

Neural net model

Used model is fully connected network built of two linear layers.

Deep Q-Learning (DQN)

The Deep Q-Learning is using a neural network, in our case a fully connected one, to approximate, given a state, the different Q-values for each possible action at that state. We use DQN_Trainer to train the agent with Q-value estimation. We create a loss function that compares our Q-value prediction for given state and the Q-target, which is the output of the model in this case.

See more at HuggingFace tutorial

Play around

Prepare python environment

Clone repository:

cd ~/
git clone https://github.com/filesmuggler/sneaky_snakes.git

Create and activate virtual environment in Python

cd sneaky_snakes
python3 -m venv ./snake_env
source ./snake_env/bin/activate

Install packages from requirements into the virtual environment

python3 -m pip install -r requirements.txt

Training

Run command:

python train.py --learning_rate 0.001 --max_mem 100000 --num_games 1000 --batch_size 1000 --width 400 --height 300 --tick 250 --scale 10

The model will be saved at the end of the session into the models directory.

Testing

Run command:

python test.py --model_path "./models/best_model_so_far.pt" --num_games 10 --width 400 --height 300 --tick 10 --scale 10

Bibliography

Coding

• GeeksForGeeks tutorial on creating snake game in Pygame from scratch
• FreeCodeCamp.org tutorial on turning snake into AI project

Theory

• HuggingFace guide on deep Q-Learning

Further reading

• Reinforcement Learning: An Introduction
• Artificial Intelligence: Foundations of Computational Agents
• Pytorch documentation
• Pygame documentation