Building a Neural Network to Play Snake

Update 6/16: This project has been selected by Apple as the WWDC20 Swift Student Challenge Winner.

Before we start…Check out the Code

This all started as a class project for computer science, and the mission was simple. Build a neural network that can play snake…from scratch. When I started on this, I didn’t know what I was getting into, but let’s see how this turned out.

The Basics

At its core, the idea of this project was to create a neural network (a small-scale simulation of the human brain) from scratch and to teach it to play the classic game Snake. This means essentially creating an artificial intelligence that is only aware of the Snake game, and then training it to play the game effectively.

Now for all the nitty-gritty details…I’ll break this into 3 sections:

1. Creating Snake
2. Creating the Neural Network
3. Bringing it all together

Creating Snake

Since I started from scratch, I needed to build the core snake game. I’ve attached my swift code for reference, and I will try to work through most of the concepts for the game.

The Snake game itself is built on 4 Swift files:

1. Coordinate.swift – The snake game itself lives inside a 20×20 grid. This class allows us to handle all aspects of a coordinate on that grid (X, Y, Position on screen, etc.) in one object.
2. Snake.swift – Defines and handles all aspects of the game related to the snake. These actions include:
   1. Initializing the snake with a random position and direction
   2. Updating the snake’s position (on the backend)
   3. Handling the snake’s position (Ex. ending the game)
   4. Resetting the snake once a game ends
3. Food.swift – Implements the food in the game. Including randomizing its position, and updating if eaten.
4. Game.swift – Handles the game as a whole including timing, starting/ending the game, and integrating the neural net.

This core game functionality is tied together and presented to the user via the MainScene.swift file. This provides an intuitive method of user interaction and allows for the user to provide manual input, toggle the neural network, change game speed, and other interactions that make for a great user experience.

Creating the Neural Network

The real bread and butter of this project is the implementation of the neural network. This neural network is built in pure swift and thus allows for extremely fast training and inferencing.

Much of the implementation of the neural network is built using SwiftAI; however, I have slimmed it down to only the aspects that this project can take advantage of.

The implementation of the neural network relies on 8 Swift files. Most are very simple, but I will elaborate on the more complex ones.

1. NeuralNet.swift – The actual implementation of the neural network. This is the backbone of the entire project and is responsible for integrating all aspects of the neural network. This includes:
   1. Setting Default Weights
   1. Training (via backpropagation)
   2. Inferring
2. Activation.swift – Implements the Relu and Softmax activation functions that are used for the hidden and output nodes respectively.
3. Cache.swift – Allows for future implementation of saving the neural net for long term storage
4. Dataset.swift – Defines the datasets that are used to train and validate the neural network at runtime. (Training Data, Training Labels, Validation Data, Validation Labels)
5. Error.swift – Implements the error function (Sometimes called cost function) that is used for the neural net. In my case, it is Mean Squared (Sometimes called Quadratic Cost).
6. GameFeatures.swift – Generates input features for the neural network to make inferences. These are based on the current state of the game. There are 5 features provided to the network:
   1. Is there an object in the “up” direction (0 or 1)
   2. Is there an object in the “right” direction (0 or 1)
   3. Is there an object in the “down” direction (0 or 1)
   4. Is there an object in the “left” direction (0 or 1)
   5. The angle between the snake and the food normalized to between -1 – 1
7. Structure.swift – Defines the structure of the neural network, including the number of layers, activation functions, and learning rate.
8. TrainingData.swift – Training data and validation data that I have included in the project by default. These were generated by me playing snake and recording the output.

Using the above classes, the neural network can be declared in Game.swift, trained at runtime (currently < 1sec), and implemented into the game.

I have found the best results using a structure of 5 input nodes, 3 layers of 100 input nodes, and 4 output nodes for the neural network.

Bringing It All Together

Overall the neural network in its current state is capable of achieving a high score of 17 under ideal circumstances.

There is still much room for improvement, specifically with the training data provided. A more robust set of training data could allow for the neural network to score even higher. This is currently the main limitation of this project.

If you would like to learn more, shoot me an email! Theres a link at the top of the page. Also, feel free to check out the code on Github and maybe leave a ⭐️.

🏄‍♂️