RL Baselines3 Zoo

RL Baselines3 Zoo is a training framework for Reinforcement Learning (RL) built on top of Stable Baselines3 (SB3). It provides a suite of scripts to automate the reinforcement learning workflow, including:

• Training: Scripts to train RL agents.
• Evaluation: Tools to evaluate trained agents.
• Hyperparameter Tuning: Automated tuning of agent parameters.
• Visualization: Scripts for plotting results and recording videos.
• Pre-trained Agents: A collection of tuned hyperparameters and trained agents for common environments and algorithms.

To install RL Zoo from the source repository, clone the repository and install it in editable mode.

git clone https://github.com/DLR-RM/rl-baselines3-zoo
cd rl-baselines3-zoo/
pip install -e .

To train an agent, ensure the environment is defined in the corresponding hyperparameter file (e.g., hyperparameters/algo_name.yml). You can initiate training using the rl_zoo3.train module.

To include evaluation and checkpoints, use the --eval-freq and --save-freq flags.

You can run a trained agent in an environment using the enjoy.py script.

If you are using the pre-trained agents provided in the repository, ensure you have cloned the repo recursively to include submodules: git clone --recursive https://github.com/DLR-RM/rl-baselines3-zoo.

To enjoy an agent from the pre-trained collection, specify the algorithm, environment, folder, and number of timesteps.

If you are enjoying an agent you trained yourself, you must specify the log folder (-f) and the experiment ID (--exp-id). Note that --exp-id 0 corresponds to the last experiment.

To download trained models from the Huggingface Hub, use the rl_zoo3.load_from_hub module. This allows you to pull existing agents directly into your local environment for evaluation or further training.

python -m rl_zoo3.load_from_hub --algo ppo --env CartPole-v1 -f logs/ -orga sb3

You can perform distributed optimization by using a shared database for Optuna. Specify the --study-name and the --storage path (e.g., a file or database URL) to allow multiple processes to contribute to the same study.

python train.py --algo ppo --env MountainCar-v0 -optimize --study-name test --storage logs/demo.log

RL Baselines3 Zoo supports experiment tracking (learning curves, hyperparameters, etc.) via Weights and Biases (W&B).

To enable tracking during training, use the --track flag. You can specify the project name with --wandb-project-name and add custom tags to the run using --wandb-tags.

Training is configured via hyperparameter files located in hyperparameters/algo_name.yml. To train an agent, ensure the environment is defined in the corresponding algorithm's YAML file.

Basic training command:

python train.py --algo algo_name --env env_id

To include evaluation during training (e.g., evaluating every 10,000 steps using 10 episodes on a single evaluation environment):

python train.py --algo sac --env HalfCheetahBulletEnv-v0 --eval-freq 10000 --eval-episodes 10 --n-eval-envs 1

You can use Docker to run RL Zoo in isolated environments. You can build your own images for CPU or GPU, or pull pre-built images from Docker Hub.

To run a script inside the CPU Docker image, use the provided helper script.