# Winner Takes It All: Training Performant RL Populations for Combinatorial Optimization

---

This repository contains the official JAX implementation of the paper [Winner Takes It All: Training Performant RL Populations for Combinatorial Optimization](LINK).

---

## Overview
Though applying reinforcement learning to combinatorial optimization is attractive, it is unrealistic to expect an agent to solve these (often NP-)hard problems in a single shot due to their inherent complexity.
Poppy is a method that uses a _population_ of agents with suitably diverse policies to improve the exploration of the
solution space of hard _combinatorial optimization_ problems, such as 
[TSP](https://en.wikipedia.org/wiki/Travelling_salesman_problem), 
[CVRP](https://en.wikipedia.org/wiki/Vehicle_routing_problem) or 
[Knapsack](https://en.wikipedia.org/wiki/Knapsack_problem). To this end, it uses a new RL objective to induce an unsupervised specialization targeted solely at maximizing the performance of the whole population.

## Using the Makefile

### Building and Running the Docker Container
For TPUs, You simply need to reach the `Makefile` remotely and run:
```shell
make docker_build_tpu
```
For CPUs/GPUs:
```shell
make docker_build_local
```

Before proceeding, you can export your Neptune API token to log results to Neptune:
```shell
export NEPTUNE_API_TOKEN="..."
```

Then you can run the Docker container and get into it to run your experiments:
```shell
make docker_start
make docker_enter
```


## Training
The models can be trained using the `experiments/run_training.py` script by executing:
```shell
python experiments/run_training.py <config> --neptune_name <neptune_name> 
```
To log the results to Neptune, add the `--neptune` flag. To disable Neptune logging, add the `--neptune_disable` flag.

where `config` is the path to a YAML configuration file specifying the main parameters we have modified across
executions, such as the environment or the population size. An example of such file is given below (a description of the
fields is given next to them).

```yaml
environment:
  name: tsp  # Other values: cvrp, knapsack
  params:
    #  CVRP Params (other pairs are <150, 55>, <150, 60>)
    num_nodes: 100
    norm_factor: 50

    # Knapsack Params
    num_items: 100
    total_budget: 25

    # TSP Params (other values are 125, 150)
    num_cities: 100

num_steps: 1000                 # Number of training steps
validation_freq: 500            # Every how many steps validation runs

pop_size: 1                     # Population size (e.g. 1, 4, 8, 16, 32)
num_starting_points: 20         # Number of starting points (must lower or equal than the problem size)
training_method: "all"          # What method is used to train the agents (all, best_one_with_tie_breaking, best_one_without_tie_breaking, anneal_parallel_best)
loss_annealing_steps: 0         # For how many steps the annealing from parallel to best works (only if training_method is anneal_parallel_best)
train_encoder: True             # Whether to train the encoder or keep its parameters frozen

num_problems_training: 64           # Batch size (number of problems for each training step)
minibatch_size_training: 64         # Training is run in minibatches (should be a divisor of num_problems_training)
num_problems_validation: 128        # Number of problems to run validation on (the set is randomly generated and fixed)
minibatch_size_validation: 128      # Validation is run in minibatches (should be a divisor of num_problems_validation)
use_augmentation_validation: False  # Whether to use augmentations during validation

save_checkpoint: True           # Whether to save checkpoints
save_best: False                # Whether to save a new checkpoint upon achieving best performance (o.w. save the last one)

load_checkpoint: ""             # Path to a checkpoint to be loaded (empty string -> do not load anything)
load_decoder: False             # Whether to load the decoder(s) from the checkpoint
load_optimizer: False           # Whether to load the optimizers from the checkpoint
```

#### Training Pipeline
The training of our models is divided into 3 phases, as explained in the paper. We briefly describe here which
parameters should be used in each of these phases for a (final) 4-agent population:
1. Train a single-decoder architecture (i.e. single agent) is trained from scratch.
```yaml
pop_size: 1
training_method: "all"   
train_encoder: True
load_checkpoint: ""
```
pop_size: 1                     # Population size (e.g. 1, 4, 8, 16, 32)
num_starting_points: 20         # Number of starting points (must lower or equal than the problem size)
training_method: "all"          # What method is used to train the agents (all, best_one)
train_encoder: False   

2. The parameters in the decoding heads are replicated as many times as agents we have in our new population. The whole model
is trained using the population objective
```yaml
pop_size: 4
train_best: "best_one"
train_encoder: True
load_checkpoint: "path to the checkpoint produced in Phase 1"
load_decoder: True
load_optimizer: False
```

### Testing
The trained models can be evaluated using the sets made available in the papers below:
* The evaluation sets for TSP100 and CVRP100 are due to [Kool et al.](https://github.com/wouterkool/attention-learn-to-route).
* The evaluation sets for TSP125, TSP150, CVRP125 and CVRP150 are due to [Hottung et al.](https://arxiv.org/pdf/2106.05126.pdf). 
* The evaluation sets for Knapsack have been produced by us.
They should be put in the `experiments/evaluation_data` folder. We describe two ways of evaluating the models below.

#### Greedy Evaluation
This evaluation runs the greedy policy (i.e. the action with the highest probability is selected). The trained models 
can be tested using the `experiments/run_evaluation.py` script by executing the command below:
```shell
python experiments/run_evaluation.py <env> <problem_size> <pop_size> <model_path>
```
where
* `env` - Name of the environment (`cvrp`, `knapsack` or `tsp`).
* `problem_size` - Size of the problem (number of nodes in CVRP, number of items in Knapsack and number of cities in TSP).
The problem sizes we have test sets for are: 100, 125 and 150 for CVRP and TSP, and 100 for Knapsack.
* `pop_size` - Population size with which the model was trained.
* `model_path` - Path to the model's folder (i.e. a checkpoint generated with the training code).

After a variable amount of time, the script will output the corresponding performance metric (e.g., tour length in TSP).

#### Sampling Evaluation
This evaluation runs stochastic rollouts for the best <agent, starting point> pairs. First, a greedy rollout is run on 
each starting position for each agent. Second, the best pairs are selected (as many as agents) and 200 stochastic
rollouts are run for each of them. The best performance is then returned.

The trained models can be tested using the `experiments/run_sampling.py` script by executing the command below:
```shell
python experiments/run_sampling.py <env> <problem_size> <pop_size> <model_path>
```
where the parameters are like those for the greedy evaluation.
