SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Felix Chalumeau*
InstaDeep
f.chalumeau@instadeep.com
Clément Bonnet
c.bonnet@instadeep.com
Arnu Pretorius
a.pretorius@instadeep.com
Shikha Surana*
s.surana@instadeep.com
Nathan Grinsztajn
n.grinsztajn@instadeep.com
Alexandre Laterre
a.laterre@instadeep.com
Thomas D. Barrett
t.barrett@instadeep.com
Combinatorial Optimization underpins many real-world applications and yet, designing performant algorithms to solve these complex, typically NP-hard, problems remains a significant research challenge. Reinforcement Learning (RL) provides a versatile framework for designing heuristics across a broad spectrum of problem domains. However, despite notable progress, RL has not yet supplanted industrial solvers as the go-to solution. Current approaches emphasize pre-training heuristics that construct solutions but often rely on search procedures with limited variance, such as stochastically sampling numerous solutions from a single policy or employing computationally expensive fine-tuning of the policy on individual problem instances. Building on the intuition that performant search at inference time should be anticipated during pre-training, we propose COMPASS, a novel RL approach that parameterizes a distribution of diverse and specialized policies conditioned on a continuous latent space. We evaluate COMPASS across three canonical problems - Travelling Salesman, Capacitated Vehicle Routing, and Job-Shop Scheduling - and demonstrate that our search strategy (i) outperforms state-of-the-art approaches in 9 out of 11 standard benchmarking tasks and (ii) generalizes better, surpassing all other approaches on a set of 18 procedurally transformed instance distributions.
Combinatorial Optimization (CO) has a wide range of real-world applications, from transportation (Contardo et al., 2012) and logistics (Laterre et al., 2018), to energy (Froger et al., 2016). Solving a CO problem consists of finding an ordering, labelling or subset of elements from a finite, discrete set that maximizes (or minimizes) a given objective function. As the number of feasible solutions typically grows exponentially with the problem size, CO problems are challenging (often NP-hard) to solve. As such, significant work goes into designing problem-specific heuristic approaches that,
whilst not guaranteeing the optimal answer, can often work well in practice. Reinforcement Learning (RL) offers a domain-agnostic framework to learn heuristics and has been successfully applied across a range of CO tasks (Vinyals et al., 2015; Deudon et al., 2018; Mazyavkina et al., 2021).
Concretely, leading RL methods typically train a policy to incrementally construct a solution one element at a time. However, whilst most efforts have focused on improving the one-shot quality of these construction heuristics (Kool et al., 2019; Kwon et al., 2020; Grinsztajn et al., 2022), it intuitively appears impractical to reliably produce the optimal solution to NP-hard problems within a single construction attempt. Consequently, competitive performance has to rely on combining a pre-trained policy with an additional search procedure. Nevertheless, this crucial aspect is often implemented using simple procedures such as stochastic sampling (Kool et al., 2019; Kwon et al., 2020; Grinsztajn et al., 2022), beam search (Steinbiss et al., 1994) or Monte Carlo Tree Search (MCTS) (Browne et al., 2012). An alternative approach, representing the current state-of-the-art for search-based RL (Bello et al., 2016; Hottung et al., 2022), is to actively re-train the heuristic on each new problem instance; however, this comes with clear computational and practical limitations. Strikingly, neither of these approaches pre-trains the policy in a way that could enable a fast and efficient inference time search: rather current approaches typically completely decouple both. The absence of an efficient search strategy is even more detrimental when the test instances are out of the distribution (OOD) used to train the policy, as this may cause a large difference between the learned policy and the policy leading to the optimal solution.
In this work, we aim to overcome the current limitations of search strategies used in RL when applied to CO problems. Our approach is to learn a latent space of diverse policies that can be explored at inference time in order to find the best-performing strategies for a given instance. This updates the current paradigm by enabling sampling from a policy space at inference time rather than constantly sampling the same policy (or set of policies) with stochasticity. We introduce COMPASS - COMbinatorial optimization with Policy Adaptation using Latent Space Search. COMPASS effectively creates an infinite set of diverse solvers by using a single conditioned policy and sampling the conditions from a continuous latent space. The training process encourages subareas of the latent space to specialize to sub-distributions of instances and this diversity is used at inference time to solve newly encountered instances.
We evaluate COMPASS on three popular CO problems: Travelling Salesman Problem (TSP), Capacitated Vehicle Routing Problem (CVRP) and Job-Shop Scheduling Problem (JSSP). After training on a distribution of fixed-sized instances for each problem, we evaluate our method on both in- and out-of-distribution test sets. We find that simple search strategies requiring no re-training provide both rapid and sustained improvement of the instance-specific policy, with COMPASS establishing a new state-of-the-art across all problems in this setting. Thanks to the diversity provided by its latent space, COMPASS achieves high performance even without a search budget and achieves comparable or better results than current leading few-shot methods.
Concretely, our work makes the following contributions: (i) We introduce COMPASS which leverages a latent space of diverse and specialized policies to effectively solve CO problems. (ii) We show that COMPASS allows for the efficient adaptation of instance-specific policies without re-training or sacrificing zero-shot performance. (iii) Experimentally, our approach is found to represent a new state-of-the-art for RL-based CO methods across all our considered problem types, achieving superior performance in 27 out of 29 tasks. (iv) We release fast and performant implementations of our method and its main competitors, written in JAX. We also provide all of our test sets including our procedurally transformed problem instances for easier comparison in future work.
**Construction methods for CO** Construction approaches in RL for CO incrementally build a solution by selecting one element at a time. After Hopfield and Tank (1985) first applied neural networks to TSP, Bello et al. (2016) extended these efforts by proposing to learn heuristics with RL using a Pointer Network (Vinyals et al., 2015) combined with an actor-critic framework. This approach was extended by Deudon et al. (2018) who added an attention-based city encoder, which was subsequently further extended by Kool et al. (2019) to use a general transformer architecture (Vaswani et al., 2017). The transformer has since become the standard model for a range of CO problems and is also used in this work. Kim et al. (2022) builds on Kool et al. (2019) by leveraging
symmetries of routing problems during training. Even though the majority of these construction approaches have focused on routing problems, numerous works have also tackled other classes of CO problems, especially on graphs, like Maximum Cut (Dai et al., 2017; Barrett et al., 2020), or Job Shop Scheduling Problem (JSSP), for which Zhang et al. (2020) proposed a Graph Neural Network (GNN) approach. A broader scope of (non-construction) approaches can be found in Appendix K.
Improving solutions at inference time As it is unlikely that the first solution generated by a construction heuristic is optimal, a popular approach consists in sampling various trajectories during inference for the same CO problem. POMO (Kwon et al., 2020) uses one policy rolled out on several versions of the same problem, while considering different starting points or symmetries, to create diverse trajectories and select the best one. Choo et al. (2022) proposes an efficient search guided by simulations, but cannot take advantage of a large inference budget by itself. EAS (Hottung et al., 2022) adds on POMO by fine-tuning a subset of the model parameters at inference time using gradient descent. However, the new solutions are biased toward the underlying pre-trained policy and can easily be stuck in local optima. Instead, MDAM (Xin et al., 2021) and Poppy (Grinsztajn et al., 2022) employ a population of agents, all of which are simultaneously rolled out at inference time. MDAM trains these policies to select different initial actions, whereas Poppy utilizes a loss function designed to specialize each policy on specific subsets of the problem distribution. Despite demonstrating promising performance, these approaches are constrained by the number of policies used during training, which remains fixed. Such a limitation quickly diminishes the benefits of additional solution candidates sampled from the population. Our method COMPASS uses the same loss as Poppy, but, unlike their approach and that of MDAM, COMPASS is not bound to a specific number of specialized policies. Moreover, its latent space makes it possible to add additional search mechanisms over the policy space, ensuring better solutions over time. CVAE-Opt (Hottung et al., 2021), akin to our method, uses a latent space for solving routing problems, however, it has several differences. First, COMPASS is trained end-to-end with RL, hence does not necessitate pre-solved instances. Second, CVAE-Opt requires training an additional recurrent encoder for (instance, solution) pairs, whereas COMPASS uses the latent space to encode a distribution of complementary policies and can be easily applied to pre-train models. Overall, COMPASS significantly outperforms CVAE-Opt while having shorter runtime.
**Formulation** The goal of a CO problem is to find the optimal labeling of a set of discrete variables that satisfies the problem's constraints. In RL, a CO problem can be formulated as a Markov Decision Process (MDP) defined by  $M = (S, A, R, T, \gamma, H)$ . This includes the state space  $S$  with states  $s_i \in S$ , action space  $A$  with actions  $a_i \in A$ , reward function  $R: S \times A \to R$ , transition function  $T(s_{i+1}|s_i, a_i)$ , discount factor  $\gamma \in [0, 1]$ , and horizon  $H$  which denotes the episode duration. The state of a problem instance is represented as the (partial) trajectory or set of actions taken in the instance, and the next state  $s_{t+1}$  is determined by applying the chosen action  $a_t$  to the current state  $s_t$ . An agent is introduced in the MDP to interact with the CO problem and find solutions by learning a policy  $\pi: S \to A$ . The policy is trained to maximize the expected sum of discounted rewards to find the optimal solution, and this is formalized as the following learning objective:
$$
\pi^ {*} = \operatorname {a r g m a x} \mathbb {E} \left[ \sum_ {t = 0} ^ {H} \gamma^ {t} R \left(s _ {t}, a _ {t}\right) \right].
$$
3.2 COMPASS
Recall our intuition that no single policy will reliably be able to solve all instances of an NP-hard CO problem in a single inference pass. Two primary approaches to address this are the inclusion of inference time search and the deployment of a diverse set of policies to increase the chance of a near-optimal strategy being deployed. This work aims to unify and extend these approaches by training an infinitely large set of diverse and specialized policies that can subsequently be searched at inference time.
To achieve this, we propose that a single set of policy parameters condition not just on the current observation, but also on samples drawn from a continuous latent space. The training objective
Figure 1: Our method COMPASS is composed of the following two phases. A. Training - the latent space is sampled to generate vectors that the policy can condition upon. The conditioned policies are then evaluated and only the best one is trained to create specialization within the latent space. B. Inference - at inference time the latent space is searched through an evolution strategy to exploit regions with high-performing policies for each instance.
then encourages this latent space of policies to be diverse (generate a wide range of behaviors) and specialized (these behaviors are optimized for different types of problem instances from the training distribution). This latent space can then be efficiently searched during inference to find the most performant policy for a given problem instance. In this section, we describe in detail the realization of this approach, which we call COMPASS (COMbinatorial optimization with Policy Adaptation using Latent Space Search). In Fig. 1, we provide an illustrated overview of COMPASS.
Our approach offers several key advantages over traditional techniques. Compared to methods that directly train multiple, uniquely parameterized policies (Xin et al., 2021; Grinsztajn et al., 2022), training a single conditional policy can, in principle, provide a continuous distribution of an infinite number of policies. Moreover, our approach mitigates the significant training and memory overheads associated with training a population of agents. Compared to methods that rely on brute-force sampling (Kool et al., 2019; Kwon et al., 2020; Grinsztajn et al., 2022) or expensive fine-tuning (Hottung et al., 2022), our training process produces a structured latent space (where similar policies are found near to each other) that permits principled search during inference.
Latent space The latent space defines the set of policies that our model can condition itself upon. Importantly, we do not learn the distribution of this space, but rather select a prior distribution over the space from which we sample during training. In practice, we use a latent space with 16 dimensions bounded between -1 and 1, and use a uniform sampling prior.
Architecture COMPASS is agnostic to the network architecture used, so long as the resulting policy is, in some way, conditioned on the vector sampled from the latent space. This can be achieved in numerous ways, from directly concatenating the vector to the input observation to conditioning keys, queries, and values in the self-attention models commonly used for CO. We refer to Appendix D for further details about the architectures used in this work and how the latent vector is used to condition them. Whilst it is possible to train COMPASS from scratch, we found that it was simple and efficient to adapt pre-trained single-policy models. To do this, we zero-initialize any additional weights corresponding to the sample latent vector such that it has no impact at the start of training. In practice, we adapt a single-agent architecture designed for few-shot inference in all of our problem settings; POMO (Kwon et al., 2020) for TSP and CVRP, and a similar architecture taken from Jumanji (Bonnet et al., 2023) for JSSP (we also considered the current SOTA model L2D (Zhang et al., 2020), however, we found that the model from Jumanji already outperformed this approach). Full network details can be found in Appendices D.1 (TSP & CVRP) and D.2 (JSSP).
**Training** The training procedure aims to specialize subareas of the latent space to sub-distributions of problems by training the policy solely on latent vectors that achieve the best performance on a given problem. At each training step, we uniformly sample a set of  $N$  vectors from the latent space and condition the policy on each vector resulting in  $N$  conditioned policies. After evaluating each policy
on the problem instance, we train the best policy (i.e. the policy conditioned on the best-performing latent vector) on the instance. The model is updated using the gradient of our objective as given by
$$
\nabla_ {\theta} J _ {\text {c o m p a s s}} = \mathbb {E} _ {\rho \sim \mathcal {D}} \mathbb {E} _ {z _ {1}, \dots , z _ {N} \sim \mathcal {P} _ {z}} \mathbb {E} _ {\tau_ {i} \sim \pi_ {\theta (\cdot | z _ {i})}} [ \nabla_ {\theta} \log \pi_ {\theta} (\tau_ {i ^ {\star}} | z _ {i ^ {\star}}) R _ {i ^ {\star}} - \mathcal {B} _ {\rho , \theta} ], \tag {1}
$$
where  $\mathcal{D}$  is the data distribution,  $\mathcal{P}_z$  the latent space,  $z_{i}$  a latent vector,  $\pi_{\theta}$  the conditioned policy,  $\tau_{i}$  the trajectory generated by policy  $\pi_{\theta}$  conditioned on vector  $z_{i}$  and has the corresponding reward  $R_{i}$ ,  $i^{\star}$  is the index of the best performing latent vector (in the sampled set) and is expressed as  $i^{\star} = \arg \max_{i\in [1,N]}R(\tau_i)$ , and lastly,  $\mathcal{B}_{\rho ,\theta}$  is the baseline, inspired by Kwon et al. (2020). Full details of the algorithmic procedure can be found in Appendix F. Notably, our work is the first to create a specialized and diverse set of policies represented by a continuous latent space by only training the best-performing vector for each problem instance.
A key training hyperparameter is the number of condition vectors sampled during evaluation. More conditioned policies results in an increased certainty that the best-performing vector in the sampled set of conditions is the best-performing vector in the latent space. Therefore, increasing the number of sampled conditions increases the likelihood of training the true best latent vector for the given problem instance, rather than a potentially suboptimal vector. More details (including training times and environment steps) are reported in Appendix F.
Inference-time search Given the latent space of diverse, specialized policies obtained by training COMPASS, at inference time, we apply a principled search procedure to find the most performant strategies. Our desired properties for a search procedure are that it should be simple, capable of rapid adaptation and robust to local optima. As such, evolutionary strategies are an appropriate approach. Specifically, we use Covariance Matrix Adaptation (CMA-ES, (Hansen and Ostermeier, 2001)). CMA-ES uses a multivariate normal distribution to sample vectors and iteratively updates the distribution's mean to increase the expected performance of sampled vectors (i.e. the quality of the solution found by the policy corresponding to each vector). The covariance is also adapted over time, either for exploration (high values, broad sampling) or exploitation (small values, focused sampling).
For a given problem instance, there may be multiple high-performance policies, therefore we use several independent CMA-ES components in parallel. To ensure that those components explore distinct areas of the space (or at least, take different paths), we compute a Voronoi Tessellation (Du et al., 1999) of the latent space and use the corresponding centroids to initialize the means of the CMA-ES components. This method proves to be robust, easy to tune, and fast, and requires low memory and computation budget, making it the perfect candidate for efficient adaptation at inference time. In our experimental section (4.3), we present an analysis of our latent space and how it is explored by CMA-ES. Details and considered alternatives can be found in Appendix E.4.
We evaluate our method on three problems - Travelling Salesman (TSP), Capacitated Vehicle Routing (CVRP), and Job Shop Scheduling (JSSP) - widely used to assess RL-based methods for CO (Deudon et al., 2018; Kool et al., 2019; Grinsztajn et al., 2022; Hottung et al., 2022). In Section 4.1, we evaluate COMPASS in the standard setting used by other methods from the literature and report results on each problem type. In Section 4.2, we assess the robustness of methods by evaluating them on instances of TSP and CVRP that are procedurally transformed using the approach developed by Bossek et al. (2019). In Section 4.3, we analyze the methods' search strategies; in particular, we provide insights about COMPASS' latent space and how it is navigated by CMA-ES at inference time. Figure 2 provides a radar plot overview of our aggregated experimental results across six performance categories of interest: (1) in distribution instances, OOD instances with different levels of distribution shift (2) small, (3) medium and (4) large, (5) large instance sizes as well as (6) few-shot performance. Our results highlight the strengths and weaknesses of each approach and in particular, the versatility and superiority of COMPASS.
Baselines We compare COMPASS to a suite of leading RL methods and industrial solvers. Across all problems we provide baselines for EAS (Hottung et al., 2022); the current SOTA active-search RL method that fine-tunes the policy on each problem instance, and Poppy (Grinsztajn et al., 2022); the current SOTA active-search RL method that stochastically samples from a fixed population of
pre-trained solvers. For routing problems (TSP and CVRP), we also provide results for POMO (Kwon et al., 2020); the leading single-agent, one-shot architecture on which EAS and Poppy are built, and LKH (Helsgaun, 2017); a leading industrial solver. We also report results of TSP-specific industrial solver Concorde (Applegate et al., 2006). For JSSP, we provide results for L2D (Zhang et al., 2020); the leading single-agent, one-shot architecture. We also provide results for the attention-based model proposed in Jumanji (Bonnet et al., 2023) that proved to outperform L2D. Finally, we report the results of Google OR-Tools (Perron and Furnon, 2019); the reference industrial solver for JSSP.
Training As our method is capable of adopting initial parameters from a pre-trained model, we re-use publicly available checkpoints of POMO (details in Appendix H) as the starting point for COMPASS on TSP and CVRP. For JSSP, we found attention-based model from Bonnet et al. (2023) outperforms L2D and hence choose it to be the reference single-agent architecture. We train the model and use the same trained checkpoints for all methods. We then train COMPASS until convergence on the same training distribution as that used to train the initial checkpoint. For TSP and CVRP these are problem instances with 100 locations uniformly sampled within a unit square. For JSSP, we use the same training distribution used in EAS, which is an instance with 10 jobs and machines, and a maximum possible duration of 98. A single set of hyperparameters is used across all problems, with full training details provided in Appendix G.
Inference When evaluating active-search performance, each method is given a fixed budget of 1600 attempts - similar to Hottung et al. (2022); Grinsztajn et al. (2022) -, where each attempt consists of one trajectory per possible starting point. This approach is used to enable direct comparison to POMO and EAS which use rollouts from all starting points at each step. For the main results on TSP and CVRP, we do not use the "augmentation trick"; where the same problem is solved multiple times by rotating the coordinate frame to make it appear different and thus generate additional diverse trajectories. This trick was used in a few baselines from prior work, however, we refrain from using it in the main results of this work for two reasons: (1) it is a domain-specific trick mainly applicable to routing problems and (2) it significantly increases the required computational budget. We nevertheless provide some results in both settings to ease comparison with previous work. Overall, the trajectory budget is exactly the same as the one used in Grinsztajn et al. (2022); Hottung et al. (2022). Note that expressing the budget in terms of trajectories gives an advantage to EAS, which uses more time, memory, and computation due to the backpropagations used to update the policy during the search.