{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Model on VRPLib\n",
    "\n",
    "In this notebook, we will test the trained model's performance on the VRPLib benchmark. We will use the trained model from the previous notebook.\n",
    "\n",
    "[VRPLIB](http://vrp.galgos.inf.puc-rio.br/index.php/en/) is a collection of instances related to the CVRP, which is a classic optimization challenge in the field of logistics and transportation. \n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/ai4co/rl4co/blob/main/examples/datasets/2-test-on-cvrplib.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>\n",
    "\n",
    "## Before we start\n",
    "\n",
    "To use the VRPLib, we strongly recomment to use the Python `vrplib` tool:\n",
    "\n",
    "[VRPLib](https://github.com/leonlan/VRPLIB) is a Python package for working with Vehicle Routing Problem (VRP) instances. This tool can help us easily load the VRPLib instances and visualize the results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "Uncomment the following line to install the package from PyPI. Remember to choose a GPU runtime for faster training!\n",
    "\n",
    "> Note: You may need to restart the runtime in Colab after this\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install rl4co[graph] # include torch-geometric\n",
    "\n",
    "## NOTE: to install latest version from Github (may be unstable) install from source instead:\n",
    "# !pip install git+https://github.com/ai4co/rl4co.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the `tsplib95` package\n",
    "# !pip install vrplib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import re\n",
    "import torch\n",
    "import vrplib\n",
    "\n",
    "from rl4co.envs import TSPEnv, CVRPEnv\n",
    "from rl4co.models.zoo.am import AttentionModel\n",
    "from rl4co.utils.trainer import RL4COTrainer\n",
    "from rl4co.utils.decoding import get_log_likelihood\n",
    "from rl4co.models.zoo import EAS, EASLay, EASEmb, ActiveSearch\n",
    "\n",
    "from tqdm import tqdm\n",
    "from math import ceil\n",
    "from einops import repeat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load a trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cbhua/miniconda3/envs/rl4co-user/lib/python3.10/site-packages/lightning/pytorch/utilities/parsing.py:198: 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'])`.\n",
      "/home/cbhua/miniconda3/envs/rl4co-user/lib/python3.10/site-packages/lightning/pytorch/utilities/parsing.py:198: 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'])`.\n",
      "/home/cbhua/miniconda3/envs/rl4co-user/lib/python3.10/site-packages/lightning/pytorch/core/saving.py:177: Found keys that are not in the model state dict but in the checkpoint: ['baseline.baseline.model.encoder.init_embedding.init_embed.weight', 'baseline.baseline.model.encoder.init_embedding.init_embed.bias', 'baseline.baseline.model.encoder.init_embedding.init_embed_depot.weight', 'baseline.baseline.model.encoder.init_embedding.init_embed_depot.bias', 'baseline.baseline.model.encoder.net.layers.0.0.module.Wqkv.weight', 'baseline.baseline.model.encoder.net.layers.0.0.module.Wqkv.bias', 'baseline.baseline.model.encoder.net.layers.0.0.module.out_proj.weight', 'baseline.baseline.model.encoder.net.layers.0.0.module.out_proj.bias', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.0.1.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.0.2.module.0.weight', 'baseline.baseline.model.encoder.net.layers.0.2.module.0.bias', 'baseline.baseline.model.encoder.net.layers.0.2.module.2.weight', 'baseline.baseline.model.encoder.net.layers.0.2.module.2.bias', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.0.3.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.1.0.module.Wqkv.weight', 'baseline.baseline.model.encoder.net.layers.1.0.module.Wqkv.bias', 'baseline.baseline.model.encoder.net.layers.1.0.module.out_proj.weight', 'baseline.baseline.model.encoder.net.layers.1.0.module.out_proj.bias', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.1.1.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.1.2.module.0.weight', 'baseline.baseline.model.encoder.net.layers.1.2.module.0.bias', 'baseline.baseline.model.encoder.net.layers.1.2.module.2.weight', 'baseline.baseline.model.encoder.net.layers.1.2.module.2.bias', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.1.3.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.2.0.module.Wqkv.weight', 'baseline.baseline.model.encoder.net.layers.2.0.module.Wqkv.bias', 'baseline.baseline.model.encoder.net.layers.2.0.module.out_proj.weight', 'baseline.baseline.model.encoder.net.layers.2.0.module.out_proj.bias', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.2.1.normalizer.num_batches_tracked', 'baseline.baseline.model.encoder.net.layers.2.2.module.0.weight', 'baseline.baseline.model.encoder.net.layers.2.2.module.0.bias', 'baseline.baseline.model.encoder.net.layers.2.2.module.2.weight', 'baseline.baseline.model.encoder.net.layers.2.2.module.2.bias', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.weight', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.bias', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.running_mean', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.running_var', 'baseline.baseline.model.encoder.net.layers.2.3.normalizer.num_batches_tracked', 'baseline.baseline.model.decoder.context_embedding.project_context.weight', 'baseline.baseline.model.decoder.project_node_embeddings.weight', 'baseline.baseline.model.decoder.project_fixed_context.weight', 'baseline.baseline.model.decoder.logit_attention.project_out.weight']\n"
     ]
    }
   ],
   "source": [
    "# Load from checkpoint; alternatively, simply instantiate a new model\n",
    "# Note the model is trained for CVRP problem\n",
    "checkpoint_path = \"../cvrp-20.ckpt\" # modify the path to your checkpoint file\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# load checkpoint\n",
    "# checkpoint = torch.load(checkpoint_path)\n",
    "\n",
    "lit_model = AttentionModel.load_from_checkpoint(checkpoint_path, load_baseline=False)\n",
    "policy, env = lit_model.policy, lit_model.env\n",
    "policy = policy.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download vrp problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/37 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 37/37 [00:42<00:00,  1.16s/it]\n"
     ]
    }
   ],
   "source": [
    "problem_names = vrplib.list_names(low=50, high=200, vrp_type='cvrp') \n",
    "\n",
    "instances = [] # Collect Set A, B, E, F, M datasets\n",
    "for name in problem_names:\n",
    "    if 'A' in name:\n",
    "        instances.append(name)\n",
    "    elif 'B' in name:\n",
    "        instances.append(name)\n",
    "    elif 'E' in name:\n",
    "        instances.append(name)\n",
    "    elif 'F' in name:\n",
    "        instances.append(name)\n",
    "    elif 'M' in name and 'CMT' not in name:\n",
    "        instances.append(name)\n",
    "\n",
    "# Modify the path you want to save \n",
    "# Note: we don't have to create this folder in advance\n",
    "path_to_save = './vrplib/' \n",
    "\n",
    "try:\n",
    "    os.makedirs(path_to_save)\n",
    "    for instance in tqdm(instances):\n",
    "        vrplib.download_instance(instance, path_to_save)\n",
    "        vrplib.download_solution(instance, path_to_save)\n",
    "except: # already exist\n",
    "    pass "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utils function\n",
    "def normalize_coord(coord:torch.Tensor) -> torch.Tensor:\n",
    "    x, y = coord[:, 0], coord[:, 1]\n",
    "    x_min, x_max = x.min(), x.max()\n",
    "    y_min, y_max = y.min(), y.max()\n",
    "    \n",
    "    x_scaled = (x - x_min) / (x_max - x_min) \n",
    "    y_scaled = (y - y_min) / (y_max - y_min)\n",
    "    coord_scaled = torch.stack([x_scaled, y_scaled], dim=1)\n",
    "    return coord_scaled "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the greedy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem: A-n53-k7        Cost: 1371       Optimal Cost: 1010      \t Gap: 35.74%\n",
      "Problem: A-n54-k7        Cost: 1426       Optimal Cost: 1167      \t Gap: 22.19%\n",
      "Problem: A-n55-k9        Cost: 1333       Optimal Cost: 1073      \t Gap: 24.23%\n",
      "Problem: A-n60-k9        Cost: 1728       Optimal Cost: 1354      \t Gap: 27.62%\n",
      "Problem: A-n61-k9        Cost: 1297       Optimal Cost: 1034      \t Gap: 25.44%\n",
      "Problem: A-n62-k8        Cost: 1818       Optimal Cost: 1288      \t Gap: 41.15%\n",
      "Problem: A-n63-k9        Cost: 2166       Optimal Cost: 1616      \t Gap: 34.03%\n",
      "Problem: A-n63-k10       Cost: 1698       Optimal Cost: 1314      \t Gap: 29.22%\n",
      "Problem: A-n64-k9        Cost: 1805       Optimal Cost: 1401      \t Gap: 28.84%\n",
      "Problem: A-n65-k9        Cost: 1592       Optimal Cost: 1174      \t Gap: 35.60%\n",
      "Problem: A-n69-k9        Cost: 1641       Optimal Cost: 1159      \t Gap: 41.59%\n",
      "Problem: A-n80-k10       Cost: 2230       Optimal Cost: 1763      \t Gap: 26.49%\n",
      "Problem: B-n51-k7        Cost: 1270       Optimal Cost: 1032      \t Gap: 23.06%\n",
      "Problem: B-n52-k7        Cost: 994        Optimal Cost: 747       \t Gap: 33.07%\n",
      "Problem: B-n56-k7        Cost: 931        Optimal Cost: 707       \t Gap: 31.68%\n",
      "Problem: B-n57-k7        Cost: 1422       Optimal Cost: 1153      \t Gap: 23.33%\n",
      "Problem: B-n57-k9        Cost: 1889       Optimal Cost: 1598      \t Gap: 18.21%\n",
      "Problem: B-n63-k10       Cost: 1807       Optimal Cost: 1496      \t Gap: 20.79%\n",
      "Problem: B-n64-k9        Cost: 1150       Optimal Cost: 861       \t Gap: 33.57%\n",
      "Problem: B-n66-k9        Cost: 1746       Optimal Cost: 1316      \t Gap: 32.67%\n",
      "Problem: B-n67-k10       Cost: 1368       Optimal Cost: 1032      \t Gap: 32.56%\n",
      "Problem: B-n68-k9        Cost: 1737       Optimal Cost: 1272      \t Gap: 36.56%\n",
      "Problem: B-n78-k10       Cost: 1706       Optimal Cost: 1221      \t Gap: 39.72%\n",
      "Problem: E-n51-k5        Cost: 690        Optimal Cost: 521       \t Gap: 32.44%\n",
      "Problem: E-n76-k7        Cost: 1019       Optimal Cost: 682       \t Gap: 49.41%\n",
      "Problem: E-n76-k8        Cost: 1031       Optimal Cost: 735       \t Gap: 40.27%\n",
      "Problem: E-n76-k10       Cost: 1156       Optimal Cost: 830       \t Gap: 39.28%\n",
      "Problem: E-n76-k14       Cost: 1335       Optimal Cost: 1021      \t Gap: 30.75%\n",
      "Problem: E-n101-k8       Cost: 1265       Optimal Cost: 815       \t Gap: 55.21%\n",
      "Problem: E-n101-k14      Cost: 1567       Optimal Cost: 1067      \t Gap: 46.86%\n",
      "Problem: F-n72-k4        Cost: 425        Optimal Cost: 237       \t Gap: 79.32%\n",
      "Problem: F-n135-k7       Cost: 4219       Optimal Cost: 1162      \t Gap: 263.08%\n",
      "Problem: M-n101-k10      Cost: 1388       Optimal Cost: 820       \t Gap: 69.27%\n",
      "Problem: M-n121-k7       Cost: 1746       Optimal Cost: 1034      \t Gap: 68.86%\n",
      "Problem: M-n151-k12      Cost: 1906       Optimal Cost: 1015      \t Gap: 87.78%\n",
      "Problem: M-n200-k16      Cost: 2509       Optimal Cost: 1274      \t Gap: 96.94%\n",
      "Problem: M-n200-k17      Cost: 2339       Optimal Cost: 1275      \t Gap: 83.45%\n"
     ]
    }
   ],
   "source": [
    "for instance in instances:\n",
    "    problem = vrplib.read_instance(os.path.join(path_to_save, instance+'.vrp'))\n",
    "\n",
    "    coords = torch.tensor(problem['node_coord']).float()\n",
    "    coords_norm = normalize_coord(coords)\n",
    "    demand = torch.tensor(problem['demand'][1:]).float()\n",
    "    capacity = problem['capacity']\n",
    "    n = coords.shape[0]\n",
    "\n",
    "    # Prepare the tensordict\n",
    "    batch_size = 2\n",
    "    td = env.reset(batch_size=(batch_size,)).to(device)\n",
    "    td['locs'] = repeat(coords_norm, 'n d -> b n d', b=batch_size, d=2)\n",
    "    td['demand'] = repeat(demand, 'n -> b n', b=batch_size) / capacity\n",
    "    td['visited'] = torch.zeros((batch_size, 1, n), dtype=torch.uint8)\n",
    "    action_mask = torch.ones(batch_size, n, dtype=torch.bool)\n",
    "    action_mask[:, 0] = False\n",
    "    td['action_mask']  = action_mask\n",
    "\n",
    "    # Get the solution from the policy\n",
    "    with torch.no_grad():\n",
    "        out = policy(td.clone(), decode_type='greedy', return_actions=True)\n",
    "\n",
    "    # Calculate the cost on the original scale\n",
    "    td['locs'] = repeat(coords, 'n d -> b n d', b=batch_size, d=2)\n",
    "    neg_reward = env.get_reward(td, out['actions'])\n",
    "    cost = ceil(-1 * neg_reward[0].item())\n",
    "\n",
    "    # Load the optimal cost\n",
    "    solution = vrplib.read_solution(os.path.join(path_to_save, instance+'.sol'))\n",
    "    optimal_cost = solution['cost']\n",
    "\n",
    "    # Calculate the gap and print\n",
    "    gap = (cost - optimal_cost) / optimal_cost\n",
    "    print(f'Problem: {instance:<15} Cost: {cost:<10} Optimal Cost: {optimal_cost:<10}\\t Gap: {gap:.2%}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem: A-n53-k7        Cost: 1123       Optimal Cost: 1010      \t Gap: 11.19%\n",
      "Problem: A-n54-k7        Cost: 1305       Optimal Cost: 1167      \t Gap: 11.83%\n",
      "Problem: A-n55-k9        Cost: 1199       Optimal Cost: 1073      \t Gap: 11.74%\n",
      "Problem: A-n60-k9        Cost: 1534       Optimal Cost: 1354      \t Gap: 13.29%\n",
      "Problem: A-n61-k9        Cost: 1187       Optimal Cost: 1034      \t Gap: 14.80%\n",
      "Problem: A-n62-k8        Cost: 1474       Optimal Cost: 1288      \t Gap: 14.44%\n",
      "Problem: A-n63-k9        Cost: 1820       Optimal Cost: 1616      \t Gap: 12.62%\n",
      "Problem: A-n63-k10       Cost: 1505       Optimal Cost: 1314      \t Gap: 14.54%\n",
      "Problem: A-n64-k9        Cost: 1582       Optimal Cost: 1401      \t Gap: 12.92%\n",
      "Problem: A-n65-k9        Cost: 1332       Optimal Cost: 1174      \t Gap: 13.46%\n",
      "Problem: A-n69-k9        Cost: 1305       Optimal Cost: 1159      \t Gap: 12.60%\n",
      "Problem: A-n80-k10       Cost: 2044       Optimal Cost: 1763      \t Gap: 15.94%\n",
      "Problem: B-n51-k7        Cost: 1073       Optimal Cost: 1032      \t Gap: 3.97%\n",
      "Problem: B-n52-k7        Cost: 815        Optimal Cost: 747       \t Gap: 9.10%\n",
      "Problem: B-n56-k7        Cost: 792        Optimal Cost: 707       \t Gap: 12.02%\n",
      "Problem: B-n57-k7        Cost: 1219       Optimal Cost: 1153      \t Gap: 5.72%\n",
      "Problem: B-n57-k9        Cost: 1744       Optimal Cost: 1598      \t Gap: 9.14%\n",
      "Problem: B-n63-k10       Cost: 1611       Optimal Cost: 1496      \t Gap: 7.69%\n",
      "Problem: B-n64-k9        Cost: 931        Optimal Cost: 861       \t Gap: 8.13%\n",
      "Problem: B-n66-k9        Cost: 1427       Optimal Cost: 1316      \t Gap: 8.43%\n",
      "Problem: B-n67-k10       Cost: 1122       Optimal Cost: 1032      \t Gap: 8.72%\n",
      "Problem: B-n68-k9        Cost: 1382       Optimal Cost: 1272      \t Gap: 8.65%\n",
      "Problem: B-n78-k10       Cost: 1437       Optimal Cost: 1221      \t Gap: 17.69%\n",
      "Problem: E-n51-k5        Cost: 606        Optimal Cost: 521       \t Gap: 16.31%\n",
      "Problem: E-n76-k7        Cost: 816        Optimal Cost: 682       \t Gap: 19.65%\n",
      "Problem: E-n76-k8        Cost: 892        Optimal Cost: 735       \t Gap: 21.36%\n",
      "Problem: E-n76-k10       Cost: 943        Optimal Cost: 830       \t Gap: 13.61%\n",
      "Problem: E-n76-k14       Cost: 1160       Optimal Cost: 1021      \t Gap: 13.61%\n",
      "Problem: E-n101-k8       Cost: 1042       Optimal Cost: 815       \t Gap: 27.85%\n",
      "Problem: E-n101-k14      Cost: 1302       Optimal Cost: 1067      \t Gap: 22.02%\n",
      "Problem: F-n72-k4        Cost: 286        Optimal Cost: 237       \t Gap: 20.68%\n",
      "Problem: F-n135-k7       Cost: 1570       Optimal Cost: 1162      \t Gap: 35.11%\n",
      "Problem: M-n101-k10      Cost: 1037       Optimal Cost: 820       \t Gap: 26.46%\n",
      "Problem: M-n121-k7       Cost: 1283       Optimal Cost: 1034      \t Gap: 24.08%\n",
      "Problem: M-n151-k12      Cost: 1407       Optimal Cost: 1015      \t Gap: 38.62%\n",
      "Problem: M-n200-k16      Cost: 1811       Optimal Cost: 1274      \t Gap: 42.15%\n",
      "Problem: M-n200-k17      Cost: 1812       Optimal Cost: 1275      \t Gap: 42.12%\n"
     ]
    }
   ],
   "source": [
    "# Import augmented utils\n",
    "from rl4co.data.transforms import (\n",
    "    StateAugmentation as SymmetricStateAugmentation)\n",
    "from rl4co.utils.ops import batchify, unbatchify\n",
    "\n",
    "num_augment = 100\n",
    "augmentation = SymmetricStateAugmentation(num_augment=num_augment)\n",
    "\n",
    "for instance in instances:\n",
    "    problem = vrplib.read_instance(os.path.join(path_to_save, instance+'.vrp'))\n",
    "\n",
    "    coords = torch.tensor(problem['node_coord']).float()\n",
    "    coords_norm = normalize_coord(coords)\n",
    "    demand = torch.tensor(problem['demand'][1:]).float()\n",
    "    capacity = problem['capacity']\n",
    "    n = coords.shape[0]\n",
    "\n",
    "    # Prepare the tensordict\n",
    "    batch_size = 2\n",
    "    td = env.reset(batch_size=(batch_size,)).to(device)\n",
    "    td['locs'] = repeat(coords_norm, 'n d -> b n d', b=batch_size, d=2)\n",
    "    td['demand'] = repeat(demand, 'n -> b n', b=batch_size) / capacity\n",
    "    td['visited'] = torch.zeros((batch_size, 1, n), dtype=torch.uint8)\n",
    "    action_mask = torch.ones(batch_size, n, dtype=torch.bool)\n",
    "    action_mask[:, 0] = False\n",
    "    td['action_mask']  = action_mask\n",
    "    \n",
    "    # Augmentation\n",
    "    td = augmentation(td)\n",
    "\n",
    "    # Get the solution from the policy\n",
    "    with torch.no_grad():\n",
    "        out = policy(\n",
    "            td.clone(), decode_type='greedy', num_starts=0, return_actions=True\n",
    "        )\n",
    "\n",
    "    # Calculate the cost on the original scale\n",
    "    coords_repeat = repeat(coords, 'n d -> b n d', b=batch_size, d=2)\n",
    "    td['locs'] = batchify(coords_repeat, num_augment)\n",
    "    reward = env.get_reward(td, out['actions'])\n",
    "    reward = unbatchify(reward, num_augment)\n",
    "    cost = ceil(-1 * torch.max(reward).item())\n",
    "\n",
    "    # Load the optimal cost\n",
    "    solution = vrplib.read_solution(os.path.join(path_to_save, instance+'.sol'))\n",
    "    optimal_cost = solution['cost']\n",
    "\n",
    "    # Calculate the gap and print\n",
    "    gap = (cost - optimal_cost) / optimal_cost\n",
    "    print(f'Problem: {instance:<15} Cost: {cost:<10} Optimal Cost: {optimal_cost:<10}\\t Gap: {gap:.2%}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem: A-n53-k7        Cost: 1191       Optimal Cost: 1010      \t Gap: 17.92%\n",
      "Problem: A-n54-k7        Cost: 1328       Optimal Cost: 1167      \t Gap: 13.80%\n",
      "Problem: A-n55-k9        Cost: 1286       Optimal Cost: 1073      \t Gap: 19.85%\n",
      "Problem: A-n60-k9        Cost: 1631       Optimal Cost: 1354      \t Gap: 20.46%\n",
      "Problem: A-n61-k9        Cost: 1230       Optimal Cost: 1034      \t Gap: 18.96%\n",
      "Problem: A-n62-k8        Cost: 1505       Optimal Cost: 1288      \t Gap: 16.85%\n",
      "Problem: A-n63-k9        Cost: 1840       Optimal Cost: 1616      \t Gap: 13.86%\n",
      "Problem: A-n63-k10       Cost: 1590       Optimal Cost: 1314      \t Gap: 21.00%\n",
      "Problem: A-n64-k9        Cost: 1643       Optimal Cost: 1401      \t Gap: 17.27%\n",
      "Problem: A-n65-k9        Cost: 1381       Optimal Cost: 1174      \t Gap: 17.63%\n",
      "Problem: A-n69-k9        Cost: 1451       Optimal Cost: 1159      \t Gap: 25.19%\n",
      "Problem: A-n80-k10       Cost: 2170       Optimal Cost: 1763      \t Gap: 23.09%\n",
      "Problem: B-n51-k7        Cost: 1187       Optimal Cost: 1032      \t Gap: 15.02%\n",
      "Problem: B-n52-k7        Cost: 884        Optimal Cost: 747       \t Gap: 18.34%\n",
      "Problem: B-n56-k7        Cost: 853        Optimal Cost: 707       \t Gap: 20.65%\n",
      "Problem: B-n57-k7        Cost: 1314       Optimal Cost: 1153      \t Gap: 13.96%\n",
      "Problem: B-n57-k9        Cost: 1744       Optimal Cost: 1598      \t Gap: 9.14%\n",
      "Problem: B-n63-k10       Cost: 1698       Optimal Cost: 1496      \t Gap: 13.50%\n",
      "Problem: B-n64-k9        Cost: 1045       Optimal Cost: 861       \t Gap: 21.37%\n",
      "Problem: B-n66-k9        Cost: 1506       Optimal Cost: 1316      \t Gap: 14.44%\n",
      "Problem: B-n67-k10       Cost: 1254       Optimal Cost: 1032      \t Gap: 21.51%\n",
      "Problem: B-n68-k9        Cost: 1510       Optimal Cost: 1272      \t Gap: 18.71%\n",
      "Problem: B-n78-k10       Cost: 1514       Optimal Cost: 1221      \t Gap: 24.00%\n",
      "Problem: E-n51-k5        Cost: 613        Optimal Cost: 521       \t Gap: 17.66%\n",
      "Problem: E-n76-k7        Cost: 882        Optimal Cost: 682       \t Gap: 29.33%\n",
      "Problem: E-n76-k8        Cost: 952        Optimal Cost: 735       \t Gap: 29.52%\n",
      "Problem: E-n76-k10       Cost: 1015       Optimal Cost: 830       \t Gap: 22.29%\n",
      "Problem: E-n76-k14       Cost: 1185       Optimal Cost: 1021      \t Gap: 16.06%\n",
      "Problem: E-n101-k8       Cost: 1189       Optimal Cost: 815       \t Gap: 45.89%\n",
      "Problem: E-n101-k14      Cost: 1420       Optimal Cost: 1067      \t Gap: 33.08%\n",
      "Problem: F-n72-k4        Cost: 344        Optimal Cost: 237       \t Gap: 45.15%\n",
      "Problem: F-n135-k7       Cost: 3130       Optimal Cost: 1162      \t Gap: 169.36%\n",
      "Problem: M-n101-k10      Cost: 1221       Optimal Cost: 820       \t Gap: 48.90%\n",
      "Problem: M-n121-k7       Cost: 1538       Optimal Cost: 1034      \t Gap: 48.74%\n",
      "Problem: M-n151-k12      Cost: 1688       Optimal Cost: 1015      \t Gap: 66.31%\n",
      "Problem: M-n200-k16      Cost: 2252       Optimal Cost: 1274      \t Gap: 76.77%\n",
      "Problem: M-n200-k17      Cost: 2260       Optimal Cost: 1275      \t Gap: 77.25%\n"
     ]
    }
   ],
   "source": [
    "# Parameters for sampling\n",
    "num_samples = 100\n",
    "softmax_temp = 0.05\n",
    "\n",
    "for instance in instances:\n",
    "    problem = vrplib.read_instance(os.path.join(path_to_save, instance+'.vrp'))\n",
    "\n",
    "    coords = torch.tensor(problem['node_coord']).float()\n",
    "    coords_norm = normalize_coord(coords)\n",
    "    demand = torch.tensor(problem['demand'][1:]).float()\n",
    "    capacity = problem['capacity']\n",
    "    n = coords.shape[0]\n",
    "\n",
    "    # Prepare the tensordict\n",
    "    batch_size = 2\n",
    "    td = env.reset(batch_size=(batch_size,)).to(device)\n",
    "    td['locs'] = repeat(coords_norm, 'n d -> b n d', b=batch_size, d=2)\n",
    "    td['demand'] = repeat(demand, 'n -> b n', b=batch_size) / capacity\n",
    "    td['visited'] = torch.zeros((batch_size, 1, n), dtype=torch.uint8)\n",
    "    action_mask = torch.ones(batch_size, n, dtype=torch.bool)\n",
    "    action_mask[:, 0] = False\n",
    "    td['action_mask']  = action_mask\n",
    "    \n",
    "    # Sampling\n",
    "    td = batchify(td, num_samples)\n",
    "\n",
    "    # Get the solution from the policy\n",
    "    with torch.no_grad():\n",
    "        out = policy(\n",
    "            td.clone(), decode_type='sampling', num_starts=0, return_actions=True\n",
    "        )\n",
    "\n",
    "    # Calculate the cost on the original scale\n",
    "    coords_repeat = repeat(coords, 'n d -> b n d', b=batch_size, d=2)\n",
    "    td['locs'] = batchify(coords_repeat, num_samples)\n",
    "    reward = env.get_reward(td, out['actions'])\n",
    "    reward = unbatchify(reward, num_samples)\n",
    "    cost = ceil(-1 * torch.max(reward).item())\n",
    "\n",
    "    # Load the optimal cost\n",
    "    solution = vrplib.read_solution(os.path.join(path_to_save, instance+'.sol'))\n",
    "    optimal_cost = solution['cost']\n",
    "\n",
    "    # Calculate the gap and print\n",
    "    gap = (cost - optimal_cost) / optimal_cost\n",
    "    print(f'Problem: {instance:<15} Cost: {cost:<10} Optimal Cost: {optimal_cost:<10}\\t Gap: {gap:.2%}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rl4co-user",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}