The aim of this project is to disseminate deep reinforcement learning agents trained by a variety of algorithms, and to enable easy analysis, comparision, and visualization of them. The hope is to reduce friction for further research into understanding reinforcement learning agents. This project makes use of the excellent Lucid neural network visualization library, and integrates with the Dopamine model release.

A paper introducing this work was published at IJCAI 2019: An Atari Model Zoo for Analyzing, Visualizing, and Comparing Deep Reinforcement Learning Agents.

This software package is accompanied by a binary release of (1) frozen models trained on Atari games by a variety of deep reinforcement learning methods, and (2) cached gameplay experience of those agents in their training environments, which is hosted online.

Dependencies:

• tensorflow (with version >0.8 <2.0; we are not currently supporting TF 2.x)
• lucid (version
• matplotlib for some visualiztions
• moviepy (optional for making movies)
• gym (installed with support for Atari; optional for generating new rollouts)
• opencv-python (optional for generating new rollouts)
• tensorflow-onnx (optional for exporting to ONNX format)

To install, run setup.py install after installing dependencies.

import atari_zoo
from atari_zoo import MakeAtariModel
from pylab import *

algo = "a2c"
env = "ZaxxonNoFrameskip-v4"
run_id = 1
tag = "final"
m = MakeAtariModel(algo,env,run_id,tag)()

# get observations, frames, and ram state from a representative rollout
obs = m.get_observations()
frames = m.get_frames()
ram = m.get_ram()

# visualize first layer of convolutional weights
session = atari_zoo.utils.get_session()

m.load_graphdef()
m.import_graph()

conv_weights = m.get_weights(session,0)
atari_zoo.utils.visualize_conv_w(conv_weights)
show()

From the command line you can run: python -m atari_zoo.activation_movie --algo rainbow --environment PongNoFrameskip-v4 --run_id 1 --output ./pong_rainbow1_activation.mp4

Example jupyter notebooks live in the notebook directory that give further examples of how this library can be used.

A starter colab notebook enables you to check out the library without downloading and installing it.

• A tool for viewing videos of trained agents is available here; note that it is possible to link to specific videos, e.g. https://uber-research.github.io/atari-model-zoo/video.html?algo=apex&game=Seaquest&tag=final&run=2.

• A tool for viewing videos of trained agents alongside their neural activations is available here.

We trained four algorithms ourselves:

• A2C - we used the baselines package from OpenAI
• GA - we used the fast GPU implementation version released by Uber
• ES - we used the fast GPU version released by Uber
• Ape-X - we used the replication released by Uber
• IMPALA - we used the released code from DeepMind

We took trained final models from two algorithms (DQN and Rainbow) from the Dopamine model release:

• DQN - implementation here
• Rainbow - implementation here

To cite this work in publications, please use the following BibTex entry:

@inproceedings{
title = {An Atari Model Zoo for Analyzing, Visualizing, and Comparing Deep Reinforcement Learning Agents},
author = {Felipe Such, Vashish Madhavan, Rosanne Liu, Rui Wang, Pablo Castro, Yulun Li, Jiale Zhi, Ludwig Schubert, Marc G. Bellemare, Jeff Clune, Joel Lehman},
booktitle = {Proceedings of IJCAI 2019},
year = {2019},
}

For questions, comments, and suggestions, email joel.lehman@uber.com.