- Overview
- About AtariAgent
- About CerberusAgent
- How to run the code
- Class Diagram
- Train agent from scratch
- Repository Structure
- References
- Update June 2026: AtariAgent reaches maximum score
- Update June 2026: Training AtariAgent with Gymnasium + Stable-Baselines3
- Update July 2026: AtariAgent++ and PPO is all you need
In this project we've trained 3 agents, namely: BlindAgent, LidarAgent, AtariAgent.
Each of them has its own state-space representation and
all of them are trained using the DQN algoritm (more info at this link).
After identifying the best state-space representation, the CerberusAgent was built by ensembling three pre-trained agents
(3 AtariAgents). You can find more information in the report.pdf.
In the following we show the record achieved by each agent.
| CerberusAgent |
|---|
|
| BlindAgent | LidarAgent | AtariAgent |
|---|---|---|
|
|
|
I will not introduce all the agents since they're explained in the report.pdf, in this readme i will present only the
CerberusAgent. However, to introduce it, I need to present the AtariAgent.
The AtariAgent takes its name from the famous DeepMind paper. In the paper they use a CNN and the input is built by stacking the last 4 frames of the game, the AtariAgent use the same concept. Its Q-network has the following structure:
Is AtariAgent not enough to win Snake? No. Well, at least I was't able to train it to do so (not true, see below :) ). By looking at the record achieved by the AtariAgent you will see that the policy learned can be splitted in two: at the beginning the snake rushes towards the food (like BlindAgent does), then it switches to a more survival policy and it starts to follow a circular path.
So we can just train 2 different AtariAgent: the first one is the one we have already trained (aka the one playing in the gif); the second one is trained using a longer snake (like 50 for example) from the beginning. In this way the second one must learn a survival strategy and by combining the two agents we should win the game.
The CerberusAgent comes from that observation. This agent is composed by 3 AtariAgents, we call them heads. Why 3 and not just 2? The short answer is: because it works better in this way. The idea behind the introduction of the second head, is related to the input distribution of the third head. However note that, by removing it, you will get similar results.
At each step of the game, a single head gives the next action to perform. The head that decides the action is selected based on the current score. The simple logic is reported in the following picture:
Hell no. The overall strategy can win the game but it is far from being a perfect strategy.
In the picture above we can see the score distribution. The bars are colored according to the head that was responsible for the snake’s death. Note that: when game starts, the lenght of the snake is 3 and the maximum is 100, so the maximum score is 97 (the score is the number of apples eaten). We observe that the agent reaches the end-game very frequently, although it is less likely to fully win the game.
First install the requirements:
pip install -r requirements.txtPrint help message with python main.py --help, the output:
usage: main.py [-h] [--agent AGENT] [--train] [--loadbuf] [--record] [--nogui]
optional arguments:
-h, --help show this help message and exit
--agent AGENT Specify agent type. Values: atari, blind, lidar or baseline.
--train If you want to train the agent.
--loadbuf If you want to load the buffer to 'device'.
--record If you want to watch the record of the agent.
--nogui If you want to deactivate the gui.
Below I report all the commands you may want to try:
# Run baseline algorithm
python main.py --agent baseline
# See replay of the record
python main.py --agent cerberus --record
python main.py --agent blind --record
python main.py --agent lidar --record
# Train agent
python main.py --agent cerberus --train
python main.py --agent blind --train
python main.py --agent lidar --train
# Create plots with stats (stored in ./output/plots/)
python plot.pyNote that: with --train you're not training a model from scratch.
It will automatically load a pre-trained model stored in ./output/models/.
Note that: the training stats are stored in ./output/csv/stats.csv.
If you stop training and then restart, the file will not be deleted, new
information will be added at the bottom. Moreover, the output model will be
stored to ./output/models/model.pth and by default a snapshot is created
only when the agent achives a new record.
When you use --train you can also add option --nogui to disable the gui (you just get the terminal output
with the info). Also you can add option --loadbuf: this will load a replay buffer stored in ./output/models/.
You can also plot the stats of the pre-trained models. For example, to plot the stats of the BlindAgent:
# Copy stats
cp ./output/csv/stats_BlindAgent.csv ./output/csv/stats.csv"
python plot.py
If you want to train an agent from scratch, you need to create a new file. The minimal code you need is:
from Agent import *
from Game import ReplaySnakeGame, Point, SnakeGame
# Define the device variable.
if torch.cuda.is_available():
device = torch.device("cuda")
print("INFO: CUDA is available. Running on GPU.")
elif torch.backends.mps.is_available():
device = torch.device("mps")
print ("INFO: MPS device found. Running on GPU")
else:
device = torch.device("cpu")
print("INFO: CUDA and MPS not available. Running on CPU.")
agent = BlindAgent( # or LidarAgent or AtariAgent
max_dataset_size = 100_000,
batch_size = 32,
lr = 0.00025,
epsilon = 1, # Starting value
decaying_epsilon = 0.995, # Only for BlindAgent
min_epsilon = 0.0001,
gamma = 0.99,
target_sync = 2000,
out_model_path = "./output/models/model.pth",
memory_path = "./output/models/memory.h5",
out_csv_path = "./output/csv/stats.csv",
device = device,
gui = False,
checkpoint_path = None,
load_buffer = False
)
agent.train()Training a BlindAgent is very fast (~200 games is enough), the other two need more time.
.
├── Agent.py # Implementations of BlindAgent, LidarAgent, AtariAgent, Cerberus
├── Baseline.py # Implemetation of the Baseline Algorithm
├── Game.py # Implementation of Snake Game
├── main.py # ... main file
├── Model.py # Neural Nets used to implement the Q-Net and DQN trainer.
├── plot.py # Used to plot stats inside './output/plots/'
├── ReplayBuffer.py # Implementation of the Replay Buffer.
└── output
├── csv # Directory where stats are stored (csv format)
├── models # Some pre-trained models
├── plots # Outputs of 'plot.py'
└── uml # Class diagram uml
Here a list some useful resources:
- The starting point for Snake and Reinforcement learning -> link
- BlindAgent was inspired by this project -> link
- LidarAgent was inspired by this project. This use a Genetic Algorithm however -> link
- Another Genetic Algorithm approach -> link
AtariAgent is enough to win at snake. I've tried different solution and I've tested many different parameters and it turns out that AtariAgent can reach the maximum score. But as you can see from the plot below, the winning rate is extremely low! Also, interestingly, the score distribution has two modes (I'm not really sure why).
What made the difference? Vectorized environments. Previously (using the code of this repo) we were using a single environment: for each training iteration, the agent was taking a single transition using its current policy. This means that we collected only a single transition from each policy. As a result, during training, the replay buffer may contain transitions generated by policies that differ significantly from the current one. This is a problem: how can we improve the current policy if we have generated only a single sample with it? This problem is amplified if the buffer is huge (in my case it was 2 million): it takes many steps to remove the old samples.
With Vectorized environments solves the problem since we collect many transitions using the same policy. Note that using too many environments leads to the opposite problem: the replay buffer becomes dominated by samples collected from the current policy, reducing the diversity of the training data (high bias).
Note: I will not push the code in this repository.
After the introduction of vectorized environments I've tried to improve the performances by implementing soft target updates and prioritized replay buffer (PER) using a sum tree, but I didn't get any significant improvements.
I've made a test using the DQN implementation of SB3. The results are much better w.r.t. my implementation: higher average score and winning rate.
The main problem of the two learned policies (using SB3 and the one learned using my implementation) is that they struggle to avoid traps. I think this is mainly related to the reward function since the agent plays very aggressively. In the following tests I'll try to investigate that.
AtariAgent++ is a modified version of AtariAgent: different neural network architecture; slightly different state space and different training algorithm (PPO from SB3 implementation). Moreover, I've made some changes to the reward function.
Below is a short demonstration of the learned policy in action. Interestingly, toward the end of the game, the learned policy closely resembles the Hamiltonian cycle. Unlike the Hamiltonian baseline, however, it takes shortcuts during the early stages of the game, allowing the snake to reach the food in fewer steps and grow more quickly!
| AtariAgent++ |
|---|
 |
The code with this implementation is available at this link.