From 7a16b64fdce55ef137fbb3cfadd7005ba0f2c5a5 Mon Sep 17 00:00:00 2001
From: Federico Berto <berto.federico2@gmail.com>
Date: Thu, 24 Oct 2024 18:55:47 +0900
Subject: [PATCH] [Example] CVRPLib update

---
 examples/datasets/2-test-on-cvrplib.ipynb | 52 +++++++++++++++++------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/examples/datasets/2-test-on-cvrplib.ipynb b/examples/datasets/2-test-on-cvrplib.ipynb
index 821664dc..112028e9 100644
--- a/examples/datasets/2-test-on-cvrplib.ipynb
+++ b/examples/datasets/2-test-on-cvrplib.ipynb
@@ -4,19 +4,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Test Model on VRPLib\n",
+    "# Testing 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",
+    "In this notebook, we will test the trained model's performance on the VRPLib benchmark.\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",
+    "[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."
+    "<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>"
    ]
   },
   {
@@ -32,11 +26,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# !pip install rl4co[graph] # include torch-geometric\n",
+    "# !pip install vrplib # for reading instance files\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"
@@ -299,7 +294,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -336,7 +331,7 @@
    "source": [
     "# Plot some instances\n",
     "env.render(tds[0], actions[0].cpu())\n",
-    "env.render(tds[-2], actions[-2].cpu())git\n",
+    "env.render(tds[-2], actions[-2].cpu())\n",
     "env.render(tds[-1], actions[-1].cpu())"
    ]
   },
@@ -344,7 +339,14 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Train"
+    "### Train\n",
+    "\n",
+    "We will train for few steps just to show the effects of training a model.\n",
+    "Alternatively, we can load the a pretrained checkpoint, e.g. with:\n",
+    "\n",
+    "```python\n",
+    "model = AttentionModel.load_from_checkpoint(checkpoint_path, load_baseline=False)\n",
+    "```\n"
    ]
   },
   {
@@ -482,6 +484,13 @@
     "trainer.fit(model)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Test trained model"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 11,
@@ -593,6 +602,21 @@
     "env.render(tds[-2], actions[-2].cpu())\n",
     "env.render(tds[-1], actions[-1].cpu())"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Great! We can see that the performance vastly improved even with just few minutes of training.\n",
+    "\n",
+    "There are several ways to improve the model's performance further, such as:\n",
+    "- Training for more steps\n",
+    "- Using a different model architecture\n",
+    "- Using a different training algorithm\n",
+    "- Using a different hyperparameters\n",
+    "- Using a different `Generator` \n",
+    "- ... and many more!"
+   ]
   }
  ],
  "metadata": {