rl4co 0.0.5

RL4CO: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmark

RL4CO

Documentation | Getting Started | Usage | Contributing | Paper | Citation

---

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

Getting started

RL4CO is now available for installation on pip!

pip install rl4co

Local install and development

If you want to develop RL4CO or access the latest builds, we recommend you to install it locally with pip in editable mode:

git clone https://github.com/kaist-silab/rl4co && cd rl4co
pip install -e .

[Optional] Automatically install PyTorch with correct CUDA version

These commands will automatically install PyTorch with the right GPU version for your system:

pip install light-the-torch
python3 -m light_the_torch install -r  --upgrade torch

Note: conda is also a good candidate for hassle-free installation of PyTorch: check out the PyTorch website for more details.

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

Usage

Train model with default configuration (AM on TSP environment):

python run.py

Change experiment

Train model with chosen experiment configuration from configs/experiment/ (e.g. tsp/am, and environment with 42 cities)

python run.py experiment=tsp/am env.num_loc=42

Disable logging

python run.py experiment=test/am logger=none '~callbacks.learning_rate_monitor'

Note that ~ is used to disable a callback that would need a logger.

Create a sweep over hyperparameters (-m for multirun)

python run.py -m experiment=tsp/am  train.optimizer.lr=1e-3,1e-4,1e-5

Minimalistic Example

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

from omegaconf import DictConfig
import lightning as L
from rl4co.envs import TSPEnv
from rl4co.models.zoo.am import AttentionModel
from rl4co.tasks.rl4co import RL4COLitModule

config = DictConfig(
    {"data": {
            "train_size": 100000,
            "val_size": 10000,
            "batch_size": 512,
        },
    "optimizer": {"lr": 1e-4}}
)

# Environment, Model, and Lightning Module
env = TSPEnv(num_loc=20)
model = AttentionModel(env)
lit_module = RL4COLitModule(config, env, model)

# Trainer
trainer = L.Trainer(
    max_epochs=3, # only few epochs
    accelerator="gpu", # use GPU if available, else you can use others as "cpu"
    logger=None, # can replace with WandbLogger, TensorBoardLogger, etc.
    precision="16-mixed", # Lightning will handle faster training with mixed precision
    gradient_clip_val=1.0, # clip gradients to avoid exploding gradients
    reload_dataloaders_every_n_epochs=1, # necessary for sampling new data
)

# Fit the model
trainer.fit(lit_module)

Testing

Run tests with pytest from the root directory:

pytest tests

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

Cite us

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

@article{berto2023rl4co,
    title = {{RL4CO}: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmark},
    author={Federico Berto and Chuanbo Hua and Junyoung Park and Minsu Kim and Hyeonah Kim and Jiwoo Son and Haeyeon Kim and Joungho Kim and Jinkyoo Park},
    journal={arXiv preprint arXiv:2306.17100},
    year={2023},
    url = {https://github.com/kaist-silab/rl4co}
}