Some time ago I started experimenting with reinforcement learning to train a model that autonomously plays the game of Snake. I started by implementing a version in Pygame that’s playable by myself and can be used as a simulation environment to train the agent.

The agent observes a small window of the game state (a 7×7 grid around its head, plus directional hints toward the food) that’s used to train a CNN network outputting the next move. After enough training episodes, it gets surprisingly good at navigating without hitting walls or itself. The 7x7 grid was essential to allow for increasing lengths of the snake to avoid self collision later in game.

Another interesting yet surprisingly simple aspect of this project was to deploy the trained neural network to a web app. While the training code itself is built in Python and Pytorch, the resulting model was converted to ONNX and loaded in a compatible javascript runtime. Below you can watch the Snake agent playing the game autonomously.

Source code: github.com/dschoerk/snake