RL4CO Quickstart Notebook¶

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

In this notebook we will train the AttentionModel (AM) on the TSP environment for 20 nodes. On a GPU, this should less than 2 minutes! 🚀

Installation¶

In [6]:

## Uncomment the following line to install the package from PyPI
## You may need to restart the runtime in Colab after this
## Remember to choose a GPU runtime for faster training!

# !pip install rl4co

## Uncomment the following line to install the package from PyPI ## You may need to restart the runtime in Colab after this ## Remember to choose a GPU runtime for faster training! # !pip install rl4co

Imports¶

In [7]:

%load_ext autoreload
%autoreload 2

import torch

from rl4co.envs import TSPEnv
from rl4co.models import AttentionModelPolicy, REINFORCE
from rl4co.utils.trainer import RL4COTrainer

%load_ext autoreload %autoreload 2 import torch from rl4co.envs import TSPEnv from rl4co.models import AttentionModelPolicy, REINFORCE from rl4co.utils.trainer import RL4COTrainer

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

Environment, Policy and Model¶

Full documentation of:https://rl4.co/docs/content/api/envs/base/

• Base environment class here
• Base policy class here
• Base model class here

In [8]:

# RL4CO env based on TorchRL
env = TSPEnv(generator_params={'num_loc': 50})

# Policy: neural network, in this case with encoder-decoder architecture
policy = AttentionModelPolicy(env_name=env.name, 
                              embed_dim=128,
                              num_encoder_layers=3,
                              num_heads=8,
                            )

# RL Model: REINFORCE and greedy rollout baseline
model = REINFORCE(env, 
                    policy,
                    baseline="rollout",
                    batch_size=512,
                    train_data_size=100_000,
                    val_data_size=10_000,
                    optimizer_kwargs={"lr": 1e-4},
                    )

# RL4CO env based on TorchRL env = TSPEnv(generator_params={'num_loc': 50}) # Policy: neural network, in this case with encoder-decoder architecture policy = AttentionModelPolicy(env_name=env.name, embed_dim=128, num_encoder_layers=3, num_heads=8, ) # RL Model: REINFORCE and greedy rollout baseline model = REINFORCE(env, policy, baseline="rollout", batch_size=512, train_data_size=100_000, val_data_size=10_000, optimizer_kwargs={"lr": 1e-4}, )

Test greedy rollout with untrained model and plot¶

In [9]:

# Greedy rollouts over untrained policy
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
td_init = env.reset(batch_size=[3]).to(device)
policy = policy.to(device)
out = policy(td_init.clone(), phase="test", decode_type="greedy", return_actions=True)
actions_untrained = out['actions'].cpu().detach()
rewards_untrained = out['reward'].cpu().detach()

for i in range(3):
    print(f"Problem {i+1} | Cost: {-rewards_untrained[i]:.3f}")
    env.render(td_init[i], actions_untrained[i])

# Greedy rollouts over untrained policy device = torch.device("cuda" if torch.cuda.is_available() else "cpu") td_init = env.reset(batch_size=[3]).to(device) policy = policy.to(device) out = policy(td_init.clone(), phase="test", decode_type="greedy", return_actions=True) actions_untrained = out['actions'].cpu().detach() rewards_untrained = out['reward'].cpu().detach() for i in range(3): print(f"Problem {i+1} | Cost: {-rewards_untrained[i]:.3f}") env.render(td_init[i], actions_untrained[i])

Problem 1 | Cost: 10.648
Problem 2 | Cost: 9.375
Problem 3 | Cost: 11.713

Trainer¶

The RL4CO trainer is a wrapper around PyTorch Lightning's Trainer class which adds some functionality and more efficient defaults

In [10]:

trainer = RL4COTrainer(
    max_epochs=3,
    accelerator="gpu",
    devices=1,
    logger=None,
)

trainer = RL4COTrainer( max_epochs=3, accelerator="gpu", devices=1, logger=None, )

Using 16bit Automatic Mixed Precision (AMP)
GPU available: True (cuda), used: True
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
/home/botu/mambaforge/envs/rl4co/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py:75: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `lightning.pytorch` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default

Fit the model¶

In [11]:

trainer.fit(model)

trainer.fit(model)

val_file not set. Generating dataset instead
test_file not set. Generating dataset instead
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name     | Type                 | Params
--------------------------------------------------
0 | env      | TSPEnv               | 0     
1 | policy   | AttentionModelPolicy | 710 K 
2 | baseline | WarmupBaseline       | 710 K 
--------------------------------------------------
1.4 M     Trainable params
0         Non-trainable params
1.4 M     Total params
5.681     Total estimated model params size (MB)

Sanity Checking: |          | 0/? [00:00<?, ?it/s]

/home/botu/mambaforge/envs/rl4co/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=31` in the `DataLoader` to improve performance.
/home/botu/mambaforge/envs/rl4co/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=31` in the `DataLoader` to improve performance.

Training: |          | 0/? [00:00<?, ?it/s]

Validation: |          | 0/? [00:00<?, ?it/s]

Validation: |          | 0/? [00:00<?, ?it/s]

Validation: |          | 0/? [00:00<?, ?it/s]

`Trainer.fit` stopped: `max_epochs=3` reached.

Testing¶

In [12]:

# Greedy rollouts over trained model (same states as previous plot)
policy = model.policy.to(device)
out = policy(td_init.clone(), phase="test", decode_type="greedy", return_actions=True)
actions_trained = out['actions'].cpu().detach()

# Plotting
import matplotlib.pyplot as plt
for i, td in enumerate(td_init):
    fig, axs = plt.subplots(1,2, figsize=(11,5))
    env.render(td, actions_untrained[i], ax=axs[0]) 
    env.render(td, actions_trained[i], ax=axs[1])
    axs[0].set_title(f"Untrained | Cost = {-rewards_untrained[i].item():.3f}")
    axs[1].set_title(r"Trained $\pi_\theta$" + f"| Cost = {-out['reward'][i].item():.3f}")

# Greedy rollouts over trained model (same states as previous plot) policy = model.policy.to(device) out = policy(td_init.clone(), phase="test", decode_type="greedy", return_actions=True) actions_trained = out['actions'].cpu().detach() # Plotting import matplotlib.pyplot as plt for i, td in enumerate(td_init): fig, axs = plt.subplots(1,2, figsize=(11,5)) env.render(td, actions_untrained[i], ax=axs[0]) env.render(td, actions_trained[i], ax=axs[1]) axs[0].set_title(f"Untrained | Cost = {-rewards_untrained[i].item():.3f}") axs[1].set_title(r"Trained $\pi_\theta$" + f"| Cost = {-out['reward'][i].item():.3f}")

We can see that even after just 3 epochs, our trained AM is able to find much better solutions than the random policy! 🎉

In [13]:

# Optionally, save the checkpoint for later use (e.g. in tutorials/4-search-methods.ipynb)
trainer.save_checkpoint("tsp-quickstart.ckpt")

# Optionally, save the checkpoint for later use (e.g. in tutorials/4-search-methods.ipynb) trainer.save_checkpoint("tsp-quickstart.ckpt")