Transductive Methods¶

In this notebook, we will showcase how to use the Efficient Active Search (EAS) algorithm to find better solutions to existing problems!

Tip: in transductive RL we train (or finetune) to solve only specific ones.

Installation¶

Uncomment the following line to install the package from PyPI. Remember to choose a GPU runtime for faster training!

Note: You may need to restart the runtime in Colab after this

In [1]:

# !pip install rl4co[graph] # include torch-geometric

## NOTE: to install latest version from Github (may be unstable) install from source instead:
# !pip install git+https://github.com/ai4co/rl4co.git

# !pip install rl4co[graph] # include torch-geometric ## NOTE: to install latest version from Github (may be unstable) install from source instead: # !pip install git+https://github.com/ai4co/rl4co.git

Imports¶

In [2]:

%load_ext autoreload
%autoreload 2

import torch

from rl4co.envs import TSPEnv, CVRPEnv
from rl4co.models.zoo.am import AttentionModel
from rl4co.utils.trainer import RL4COTrainer
from rl4co.utils.decoding import get_log_likelihood
from rl4co.models.zoo import EAS, EASLay, EASEmb, ActiveSearch

import logging

%load_ext autoreload %autoreload 2 import torch from rl4co.envs import TSPEnv, CVRPEnv from rl4co.models.zoo.am import AttentionModel from rl4co.utils.trainer import RL4COTrainer from rl4co.utils.decoding import get_log_likelihood from rl4co.models.zoo import EAS, EASLay, EASEmb, ActiveSearch import logging

/home/botu/rl4co/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

/home/botu/rl4co/.venv/lib/python3.12/site-packages/torch/jit/_script.py:1491: FutureWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.
  warnings.warn(

In [3]:

# Load from checkpoint; alternatively, simply instantiate a new model
checkpoint_path = "../tsp-quickstart.ckpt" # saved by examples/1-quickstart.ipynb, trained for a few epochs only just for showing the examples

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# load checkpoint
# checkpoint = torch.load(checkpoint_path)

# weights_only=False: RL4CO checkpoints also pickle the env and policy, only load checkpoints you trust
model = AttentionModel.load_from_checkpoint(checkpoint_path, load_baseline=False, weights_only=False)
policy = model.policy.to(device)

# Load from checkpoint; alternatively, simply instantiate a new model checkpoint_path = "../tsp-quickstart.ckpt" # saved by examples/1-quickstart.ipynb, trained for a few epochs only just for showing the examples device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # load checkpoint # checkpoint = torch.load(checkpoint_path) # weights_only=False: RL4CO checkpoints also pickle the env and policy, only load checkpoints you trust model = AttentionModel.load_from_checkpoint(checkpoint_path, load_baseline=False, weights_only=False) policy = model.policy.to(device)

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/utilities/parsing.py:213: Attribute 'env' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['env'])`.
/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/utilities/parsing.py:213: Attribute 'policy' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['policy'])`.
/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/core/saving.py:197: Found keys that are not in the model state dict but in the checkpoint: ['baseline.baseline.policy.encoder.init_embedding.init_embed.weight', 'baseline.baseline.policy.encoder.init_embedding.init_embed.bias', 'baseline.baseline.policy.encoder.net.layers.0.0.module.Wqkv.weight', 'baseline.baseline.policy.encoder.net.layers.0.0.module.Wqkv.bias', 'baseline.baseline.policy.encoder.net.layers.0.0.module.out_proj.weight', 'baseline.baseline.policy.encoder.net.layers.0.0.module.out_proj.bias', 'baseline.baseline.policy.encoder.net.layers.0.1.normalizer.weight', 'baseline.baseline.policy.encoder.net.layers.0.1.normalizer.bias', 'baseline.baseline.policy.encoder.net.layers.0.1.normalizer.running_mean', 'baseline.baseline.policy.encoder.net.layers.0.1.normalizer.running_var', 'baseline.baseline.policy.encoder.net.layers.0.1.normalizer.num_batches_tracked', 'baseline.baseline.policy.encoder.net.layers.0.2.module.lins.0.weight', 'baseline.baseline.policy.encoder.net.layers.0.2.module.lins.0.bias', 'baseline.baseline.policy.encoder.net.layers.0.2.module.lins.1.weight', 'baseline.baseline.policy.encoder.net.layers.0.2.module.lins.1.bias', 'baseline.baseline.policy.encoder.net.layers.0.3.normalizer.weight', 'baseline.baseline.policy.encoder.net.layers.0.3.normalizer.bias', 'baseline.baseline.policy.encoder.net.layers.0.3.normalizer.running_mean', 'baseline.baseline.policy.encoder.net.layers.0.3.normalizer.running_var', 'baseline.baseline.policy.encoder.net.layers.0.3.normalizer.num_batches_tracked', 'baseline.baseline.policy.encoder.net.layers.1.0.module.Wqkv.weight', 'baseline.baseline.policy.encoder.net.layers.1.0.module.Wqkv.bias', 'baseline.baseline.policy.encoder.net.layers.1.0.module.out_proj.weight', 'baseline.baseline.policy.encoder.net.layers.1.0.module.out_proj.bias', 'baseline.baseline.policy.encoder.net.layers.1.1.normalizer.weight', 'baseline.baseline.policy.encoder.net.layers.1.1.normalizer.bias', 'baseline.baseline.policy.encoder.net.layers.1.1.normalizer.running_mean', 'baseline.baseline.policy.encoder.net.layers.1.1.normalizer.running_var', 'baseline.baseline.policy.encoder.net.layers.1.1.normalizer.num_batches_tracked', 'baseline.baseline.policy.encoder.net.layers.1.2.module.lins.0.weight', 'baseline.baseline.policy.encoder.net.layers.1.2.module.lins.0.bias', 'baseline.baseline.policy.encoder.net.layers.1.2.module.lins.1.weight', 'baseline.baseline.policy.encoder.net.layers.1.2.module.lins.1.bias', 'baseline.baseline.policy.encoder.net.layers.1.3.normalizer.weight', 'baseline.baseline.policy.encoder.net.layers.1.3.normalizer.bias', 'baseline.baseline.policy.encoder.net.layers.1.3.normalizer.running_mean', 'baseline.baseline.policy.encoder.net.layers.1.3.normalizer.running_var', 'baseline.baseline.policy.encoder.net.layers.1.3.normalizer.num_batches_tracked', 'baseline.baseline.policy.encoder.net.layers.2.0.module.Wqkv.weight', 'baseline.baseline.policy.encoder.net.layers.2.0.module.Wqkv.bias', 'baseline.baseline.policy.encoder.net.layers.2.0.module.out_proj.weight', 'baseline.baseline.policy.encoder.net.layers.2.0.module.out_proj.bias', 'baseline.baseline.policy.encoder.net.layers.2.1.normalizer.weight', 'baseline.baseline.policy.encoder.net.layers.2.1.normalizer.bias', 'baseline.baseline.policy.encoder.net.layers.2.1.normalizer.running_mean', 'baseline.baseline.policy.encoder.net.layers.2.1.normalizer.running_var', 'baseline.baseline.policy.encoder.net.layers.2.1.normalizer.num_batches_tracked', 'baseline.baseline.policy.encoder.net.layers.2.2.module.lins.0.weight', 'baseline.baseline.policy.encoder.net.layers.2.2.module.lins.0.bias', 'baseline.baseline.policy.encoder.net.layers.2.2.module.lins.1.weight', 'baseline.baseline.policy.encoder.net.layers.2.2.module.lins.1.bias', 'baseline.baseline.policy.encoder.net.layers.2.3.normalizer.weight', 'baseline.baseline.policy.encoder.net.layers.2.3.normalizer.bias', 'baseline.baseline.policy.encoder.net.layers.2.3.normalizer.running_mean', 'baseline.baseline.policy.encoder.net.layers.2.3.normalizer.running_var', 'baseline.baseline.policy.encoder.net.layers.2.3.normalizer.num_batches_tracked', 'baseline.baseline.policy.decoder.context_embedding.W_placeholder', 'baseline.baseline.policy.decoder.context_embedding.project_context.weight', 'baseline.baseline.policy.decoder.pointer.project_out.weight', 'baseline.baseline.policy.decoder.project_node_embeddings.weight', 'baseline.baseline.policy.decoder.project_fixed_context.weight']

In [4]:

# env = CVRPEnv(generator_params=dict(num_loc=50))
# policy = AttentionModel(env).policy.to(device)

env = TSPEnv(generator_params=dict(num_loc=50))

td = env.reset(batch_size=3).to(device)

out = policy(td)

# env = CVRPEnv(generator_params=dict(num_loc=50)) # policy = AttentionModel(env).policy.to(device) env = TSPEnv(generator_params=dict(num_loc=50)) td = env.reset(batch_size=3).to(device) out = policy(td)

In [5]:

env.render(td.cpu(), out["actions"].cpu())

env.render(td.cpu(), out["actions"].cpu())

Out[5]:

<Axes: >

EAS¶

We perform few iterations of EASLay for demonstration

In [6]:

logging.basicConfig(level=logging.DEBUG)

env.generator.num_loc = 200

dataset = env.dataset(batch_size=[2])
# eas_model = EASEmb(env, policy, dataset, batch_size=2, max_iters=20, save_path="eas_sols.pt") # alternative
eas_model = EASLay(env, policy, dataset, batch_size=2, max_iters=20, save_path="eas_sols.pt")

eas_model.setup()

logging.basicConfig(level=logging.DEBUG) env.generator.num_loc = 200 dataset = env.dataset(batch_size=[2]) # eas_model = EASEmb(env, policy, dataset, batch_size=2, max_iters=20, save_path="eas_sols.pt") # alternative eas_model = EASLay(env, policy, dataset, batch_size=2, max_iters=20, save_path="eas_sols.pt") eas_model.setup()

INFO:rl4co.models.rl.common.base:No metrics specified, using default

INFO:rl4co.models.zoo.eas.search:Setting up Efficient Active Search (EAS) with: 
- EAS Embedding: False 
- EAS Layer: True 

INFO:rl4co.data.transforms:Features not passed, defaulting to 'locs'

In [7]:

# Plot initial solution
td_dataset = next(iter(eas_model.train_dataloader()))
td_dataset = env.reset(td_dataset).to(device)
out = policy(td_dataset)

env.render(td_dataset.cpu(), out["actions"].cpu())

# Plot initial solution td_dataset = next(iter(eas_model.train_dataloader())) td_dataset = env.reset(td_dataset).to(device) out = policy(td_dataset) env.render(td_dataset.cpu(), out["actions"].cpu())

INFO:rl4co.models.common.constructive.base:Instantiated environment not provided; instantiating tsp

Out[7]:

<Axes: >

Perform search¶

In [8]:

from rl4co.utils.trainer import RL4COTrainer

trainer = RL4COTrainer(
    max_epochs=1,
    gradient_clip_val=None,
)

trainer.fit(eas_model)

from rl4co.utils.trainer import RL4COTrainer trainer = RL4COTrainer( max_epochs=1, gradient_clip_val=None, ) trainer.fit(eas_model)

WARNING:rl4co.utils.trainer:gradient_clip_val is set to None. This may lead to unstable training.

Using 16bit Automatic Mixed Precision (AMP)

GPU available: True (cuda), used: True

TPU available: False, using: 0 TPU cores

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py:76: 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
💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.

INFO:rl4co.models.zoo.eas.search:Setting up Efficient Active Search (EAS) with: 
- EAS Embedding: False 
- EAS Layer: True 

DEBUG:fsspec.local:open file: /home/botu/rl4co/examples/modeling/lightning_logs/version_0/hparams.yaml

INFO:rl4co.data.transforms:Features not passed, defaulting to 'locs'

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

INFO:rl4co.models.rl.common.base:Instantiating optimizer <Adam>

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/utilities/model_summary/model_summary.py:242: Precision 16-mixed is not supported by the model summary.  Estimated model size in MB will not be accurate. Using 32 bits instead.

┏━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name   ┃ Type                 ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ env    │ TSPEnv               │      0 │ train │     0 │
│ 1 │ policy │ AttentionModelPolicy │  710 K │ eval  │     0 │
└───┴────────┴──────────────────────┴────────┴───────┴───────┘

Trainable params: 710 K                                                                                            
Non-trainable params: 0                                                                                            
Total params: 710 K                                                                                                
Total estimated model params size (MB): 2.841                                                                      
Modules in train mode: 1                                                                                           
Modules in eval mode: 68                                                                                           
Total FLOPs: 0                                                                                                     

/home/botu/rl4co/.venv/lib/python3.12/site-packages/rich/live.py:260: UserWarning: install "ipywidgets" for Jupyter
support
  warnings.warn('install "ipywidgets" for Jupyter support')

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/utilities/_pytree.py:21: 
`isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` 
instead.

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: 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=63` in the `DataLoader` to improve performance.

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: 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=63` in the `DataLoader` to improve performance.

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py:321: The number of training
batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for 
log_every_n_steps if you want to see logs for the training epoch.

/home/botu/rl4co/.venv/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py:538: Found 68 module(s) in 
eval mode at the start of training. This may lead to unexpected behavior during training. If this is intentional, 
you can ignore this warning.

INFO:rl4co.models.rl.common.base:Instantiating optimizer <Adam>

INFO:rl4co.models.zoo.eas.search:0/20 |  Reward: -16.92 

INFO:rl4co.models.zoo.eas.search:1/20 |  Reward: -16.60 

INFO:rl4co.models.zoo.eas.search:2/20 |  Reward: -16.26 

INFO:rl4co.models.zoo.eas.search:3/20 |  Reward: -16.20 

INFO:rl4co.models.zoo.eas.search:4/20 |  Reward: -16.14 

INFO:rl4co.models.zoo.eas.search:5/20 |  Reward: -16.13 

INFO:rl4co.models.zoo.eas.search:6/20 |  Reward: -15.84 

INFO:rl4co.models.zoo.eas.search:7/20 |  Reward: -15.82 

INFO:rl4co.models.zoo.eas.search:8/20 |  Reward: -15.78 

INFO:rl4co.models.zoo.eas.search:9/20 |  Reward: -15.49 

INFO:rl4co.models.zoo.eas.search:10/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:11/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:12/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:13/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:14/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:15/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:16/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:17/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:18/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:19/20 |  Reward: -15.19 

INFO:rl4co.models.zoo.eas.search:Best reward: -15.19

INFO:rl4co.models.zoo.eas.search:Saving solutions and rewards to eas_sols.pt...

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

Load actions¶

In [9]:

# Load
actions = torch.load("eas_sols.pt")["solutions"][0].cpu()
actions = actions[:torch.count_nonzero(actions, dim=-1)] # remove trailing zeros
state = td_dataset.cpu()[0]

env.render(state, actions)

# Load actions = torch.load("eas_sols.pt")["solutions"][0].cpu() actions = actions[:torch.count_nonzero(actions, dim=-1)] # remove trailing zeros state = td_dataset.cpu()[0] env.render(state, actions)

Out[9]:

<Axes: >

Even with few iterations, the search method can clearly find better solutions than the initial ones!