Installation

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/ai4co/rl4co && cd rl4co
pip install -e .

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

Minimalistic Example

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

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

# Instantiate generator and environment
generator = TSPGenerator(num_loc=50, loc_distribution="uniform")
env = TSPEnv(generator)

# Create policy and RL model
policy = AttentionModelPolicy(env_name=env.name, num_encoder_layers=6)
model = POMO(env, policy, batch_size=64, optimizer_kwargs={"lr": 1e-4})

# Instantiate Trainer and fit
trainer = RL4COTrainer(max_epochs=10, accelerator="gpu", precision="16-mixed")
trainer.fit(model)

Tip

We recommend checking out our quickstart notebook!