This project provides a streamlined setup for training and evaluating reinforcement learning agents on Atari 2600 games. It is based on the workflow demonstrated in this YouTube tutorial and uses a collection of simple shell scripts to manage common tasks like training, evaluation, and video recording.
The core of the project leverages powerful libraries like Stable-Baselines3 and RL-Baselines3-Zoo to do the heavy lifting.
best.model.ppo.ALE.Breakout.v5.step.0.to.step.5000.mp4
- Cross-Platform Workflow: Simple Python scripts (
train.py,enjoy.pyandrecord-video.py) to abstract away complex commands and ensure compatibility across platforms. - Powered by RL Zoo: Leverages the robust framework of
rl-zoo3for training, evaluation, and hyperparameter management. - Easy to Customize: Scripts can be easily modified to train on different Atari games or with different RL algorithms.
The recommended way to install this tool is directly from PyPI:
pip install atari-reinforcement-learningThis will make the atari-rl command available in your environment.
Caution
Atari ROMs: This project uses ale-py to automatically download and install the necessary Atari ROMs during the dependency installation process. By proceeding with the installation, you are confirming that you have the legal right to use these ROMs.
To start training an agent from scratch, run the training script. This will save logs and the trained model in the logs/ directory.
atari-rl trainIf a training session was interrupted, you can resume from the last saved checkpoint.
atari-rl train --resumeOnce you have a trained model, you can watch it play the game. This script will load the best-performing model from your training history.
atari-rl enjoyTip
To watch a specific experiment, add --exp-id n, where n is the experiment number.
To save a video of your agent playing, use the recording script. The video will be saved in a videos/ folder inside the corresponding log directory.
atari-rl record-videoTip
You can also specify the experiment ID with --exp-id n and the output format with --format <svg|mp4|all>.
If you want to contribute to the project or modify the code, follow these steps to set up a development environment.
-
Clone the repository
git clone https://github.com/CosmicDNA/atari-reinforcement-learning.git cd atari-reinforcement-learning -
Create and activate a virtual environment
python -m venv .venv source .venv/bin/activate # On Windows, use: .venv\Scripts\activate
-
Install in editable mode with development dependencies This installs the package in a way that your code changes are immediately reflected, and includes tools like
ruffandpytestfor development.- Using
uv(recommended for speed):pip install uv uv pip install -e ".[dev]" - Using
pip:pip install -e ".[dev]"
- Using
This tool uses a .env file in your current working directory to manage experiment configurations. This allows you to maintain separate configurations for different experiments in different folders.
When you run any atari-rl command for the first time in a new directory, the tool will automatically create a .env file for you with default settings.
To customize an experiment, simply edit the values in this .env file. For example, to train a DQN agent on Pong, you would modify your .env file like this:
ALGO="dqn"
ENV="ALE/Pong-v5"This project stands on the shoulders of the following giants:
- RL Baselines3 Zoo: A Training Framework for Stable Baselines3 Reinforcement Learning Agents.
- Gymnasium: A Python library for developing and comparing reinforcement learning algorithms.
- PyTorch: An open-source machine learning framework that accelerates the path from research prototyping to production deployment.
- The Arcade Learning Environment: A simple framework that allows researchers and hobbyists to develop AI agents for Atari 2600 games.
- OpenCV: Pre-built CPU-only OpenCV packages for Python.