An extensive Reinforcement Learning (RL) for Combinatorial Optimization (CO) benchmark. Our goal is to provide a unified framework for RL-based CO algorithms, and to facilitate reproducible research in this field, decoupling the science from the engineering.

RL4CO is built upon: - TorchRL: official PyTorch framework for RL algorithms and vectorized environments on GPUs - TensorDict: a library to easily handle heterogeneous data such as states, actions and rewards - PyTorch Lightning: a lightweight PyTorch wrapper for high-performance AI research - Hydra: a framework for elegantly configuring complex applications

We offer flexible and efficient implementations of the following policies: - Constructive: learn to construct a solution from scratch - Autoregressive (AR): construct solutions one step at a time via a decoder - NonAutoregressive (NAR): learn to predict a heuristic, such as a heatmap, to then construct a solution - Improvement: learn to improve a pre-existing solution

We provide several utilities and modularization. For example, we modularize reusable components such as environment embeddings that can easily be swapped to solve new problems.

Getting started

RL4CO is now available for installation on pip! bash pip install rl4co

To get started, we recommend checking out our quickstart notebook or the minimalistic example below.

Install from source

This command installs the bleeding edge main version, useful for staying up-to-date with the latest developments - for instance, if a bug has been fixed since the last official release but a new release hasn’t been rolled out yet:

bash pip install -U git+https://github.com/ai4co/rl4co.git

Local install and development

We recommend local development with the blazing-fast uv package manager, for instance:

bash git clone https://github.com/ai4co/rl4co && cd rl4co uv sync --all-extras source .venv/bin/activate

This will create a new virtual environment in .venv/ and install all dependencies.

Usage

Train model with default configuration (AM on TSP environment): bash python run.py

[!TIP] You may check out this notebook to get started with Hydra!

Minimalistic Example

Here is a minimalistic example training the Attention Model with greedy rollout baseline on TSP in less than 30 lines of code:

```python from rl4co.envs.routing import TSPEnv, TSPGenerator from rl4co.models import AttentionModelPolicy, POMO from rl4co.utils import RL4COTrainer

Instantiate generator and environment

generator = TSPGenerator(numloc=50, locdistribution="uniform") env = TSPEnv(generator)

Create policy and RL model

policy = AttentionModelPolicy(envname=env.name, numencoderlayers=6) model = POMO(env, policy, batchsize=64, optimizer_kwargs={"lr": 1e-4})

Instantiate Trainer and fit

trainer = RL4COTrainer(max_epochs=10, accelerator="gpu", precision="16-mixed") trainer.fit(model) ```

Other examples can be found on our documentation!

Testing

Run tests with pytest from the root directory:

bash pytest tests

Known Bugs

You may check out the issues and discussions. We will also periodically post updates on the FAQ section.

Contributing

Have a suggestion, request, or found a bug? Feel free to open an issue or submit a pull request. If you would like to contribute, please check out our contribution guidelines here. We welcome and look forward to all contributions to RL4CO!

We are also on Slack if you have any questions or would like to discuss RL4CO with us. We are open to collaborations and would love to hear from you 🚀

Contributors

Citation

If you find RL4CO valuable for your research or applied projects:

bibtex @inproceedings{berto2025rl4co, title={{RL4CO: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmark}}, author={Federico Berto and Chuanbo Hua and Junyoung Park and Laurin Luttmann and Yining Ma and Fanchen Bu and Jiarui Wang and Haoran Ye and Minsu Kim and Sanghyeok Choi and Nayeli Gast Zepeda and Andr\'e Hottung and Jianan Zhou and Jieyi Bi and Yu Hu and Fei Liu and Hyeonah Kim and Jiwoo Son and Haeyeon Kim and Davide Angioni and Wouter Kool and Zhiguang Cao and Jie Zhang and Kijung Shin and Cathy Wu and Sungsoo Ahn and Guojie Song and Changhyun Kwon and Lin Xie and Jinkyoo Park}, booktitle={Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining}, year={2025}, url={https://github.com/ai4co/rl4co} }

Note that a previous version of RL4CO was also accepted as an oral presentation at the NeurIPS 2023 GLFrontiers Workshop. Since then, the library has greatly evolved and improved!

Join us

We invite you to join our AI4CO community, an open research group in Artificial Intelligence (AI) for Combinatorial Optimization (CO)!