# Solving the Traveling Salesman Problem with Positional Encoding

Code used for the paper "Solving the Traveling Salesman Problem with Positional Encoding".

![attention-viz](./outputs/attention-viz.png)

Tour d'horizon:

```
.
├── configs/      # Hyperparameters
├── data/         # Generate data
├── outputs/      # Trained models
├── src/          # Main code
├── evaluate.py   # Evaluate a model
├── Justfile      # Find most common commands
└── main.py       # Training a model
```

**NOTE: Because of the size limit of OpenReview, our models, data and previous state-of-the-art baselines are provided at the following address:**

https://drive.proton.me/urls/EZ1TQA8VER#9EUFquCwNSo3


## Installation

Make sure you have [`git-lfs`](https://git-lfs.com/) (used for the checkpoints
and the data), and [`uv`](https://docs.astral.sh/uv/) installed. Then:

```sh
git clone git@github.com:XXX/pens.git
cd pens
uv sync
source .venv/bin/activate
```

## Training a Model

We use [`hydra`](https://hydra.cc/). The main training configuration can be found at
`configs/default.yaml`. Here is an example of how to train PENS-AR:

```sh
python3 main.py alibi=true rope=true
```

Logging is done using [`wandb`](https://wandb.ai/). If you want to launch an offline training (thus
not requiring a wandb account), use:

```sh
python3 main.py alibi=true rope=true mode=offline
```

## Evaluation

Here's an example, evaluating PENS-A on TSP-10k:

```sh
python3 evaluate.py --experiment-dir outputs/PENS-A/ --scaling-factor --n-estimates 5 --batch-size 1 ./data/tsp-10000.npz
```

If your GPU don't have enought memory, you can use `--chunk-size 1024` to use chunked attention,
trading memory against solving runtime. Use `python3 evaluate.py --help` for more information. Have
a look at the provided `Justfile` for more examples of evaluations.

## Generate Data

**For training:**

```sh
python3 -m data.generate --parallel --num_instances 1000000 --num_nodes 100 --output_filename data/tsp-train.npz --seed 123 --reorder
```

**For evaluation:**

```sh
python3 -m data.generate --parallel --num_instances 128 --num_nodes 1000 --output_filename data/tsp-1000.npz --seed 1000
```
