SlowGuess/NeurIPS_2023_OCR huggingface dataset:
InstaDeep n.grinsztajn@instadeep.com
Daniel Furelos-Blanco
Imperial College London*
Shikha Surana
Applying reinforcement learning (RL) to combinatorial optimization problems is attractive as it removes the need for expert knowledge or pre-solved instances. However, it is unrealistic to expect an agent to solve these (often NP-)hard problems in a single shot at inference due to their inherent complexity. Thus, leading approaches often implement additional search strategies, from stochastic sampling and beam-search to explicit fine-tuning. In this paper, we argue for the benefits of learning a population of complementary policies, which can be simultaneously rolled out at inference. To this end, we introduce Poppy, a simple training procedure for populations. Instead of relying on a predefined or hand-crafted notion of diversity, Poppy induces an unsupervised specialization targeted solely at maximizing the performance of the population. We show that Poppy produces a set of complementary policies, and obtains state-of-the-art RL results on four popular NP-hard problems: traveling salesman, capacitated vehicle routing, 0-1 knapsack, and job-shop scheduling.
In recent years, machine learning (ML) approaches have overtaken algorithms that use handcrafted features and strategies across a variety of challenging tasks [Mnih et al., 2015, van den Oord et al., 2016, Silver et al., 2018, Brown et al., 2020]. In particular, solving combinatorial optimization (CO) problems – where the maxima or minima of an objective function acting on a finite set of discrete variables is sought – has attracted significant interest [Bengio et al., 2021] due to both their (often NP) hard nature and numerous practical applications across domains varying from logistics [Sbihi and Eglese, 2007] to fundamental science [Wagner, 2020].
As the search space of feasible solutions typically grows exponentially with the problem size, exact solvers can be challenging to scale; hence, CO problems are often also tackled with handcrafted heuristics using expert knowledge. Whilst a diversity of ML-based heuristics have been proposed, reinforcement learning [RL; Sutton and Barto, 2018] is a promising paradigm as it does not require pre-solved examples of these hard problems. Indeed, algorithmic improvements to RL-based CO solvers, coupled with low inference cost, and the fact that they are by design targeted at specific problem distributions, have progressively narrowed the gap with traditional solvers.
RL methods frame CO as sequential decision-making problems, and can be divided into two families [Mazyavkina et al., 2021]. First, improvement methods start from a feasible solution and iteratively
*Work completed during an internship at InstaDeep.
improve it through small modifications (actions). However, such incremental search cannot quickly access very different solutions, and requires handcrafted procedures to define a sensible action space. Second, construction methods incrementally build a solution by selecting one element at a time. In practice, it is often unrealistic for a learned heuristic to solve NP-hard problems in a single shot, therefore these methods are typically combined with search strategies, such as stochastic sampling or beam search. However, just as improvement methods are biased by the initial starting solution, construction methods are biased by the single underlying policy. Thus, a balance must be struck between the exploitation of the learned policy (which may be ill-suited for a given problem instance) and the exploration of different solutions (where the extreme case of a purely random policy will likely be highly inefficient).
In this work, we propose Poppy, a construction method that uses a population of agents with suitably diverse policies to improve the exploration of the solution space of hard CO problems. Whereas a single agent aims to perform well across the entire problem distribution, and thus has to make compromises, a population can learn a set of heuristics such that only one of these has to be performant on any given problem instance. However, realizing this intuition presents several challenges: (i) naively training a population of agents is expensive and challenging to scale, (ii) the trained population should have complementary policies that propose different solutions, and (iii) the training approach should not impose any handcrafted notion of diversity within the set of policies given the absence of clear behavioral markers aligned with performance for typical CO problems.
Challenge (i) can be addressed by sharing a large fraction of the computations across the population, specializing only lightweight policy heads to realize the diversity of agents. Moreover, this can be done on top of a pre-trained model, which we clone to produce the population. Challenges (ii) and (iii) are jointly achieved by introducing an RL objective aimed at specializing agents on distinct subsets of the problem distribution. Concretely, we derive a policy gradient method for the population-level objective, which corresponds to training only the agent which performs best on each problem. This is intuitively justified as the performance of the population on a given problem is not improved by training an agent on an instance where another agent already has better performance. Strikingly, we find that judicious application of this conceptually simple objective gives rise to a population where the diversity of policies is obtained without explicit supervision (and hence is applicable across a range of problems without modification) and essential for strong performance.
1. We motivate the use of populations for CO problems as an efficient way to explore environments that are not reliably solved by single-shot inference.
2. We derive a new training objective and present a practical training procedure that encourages performance-driven diversity (i.e. effective diversity without the use of explicit behavioral markers or other external supervision).
3. We evaluate Poppy on four CO problems: traveling salesman (TSP), capacitated vehicle routing (CVRP), 0-1 knapsack (KP), and job-shop scheduling (JSSP). In these four problems, Poppy consistently outperforms all other RL-based approaches.
ML for Combinatorial Optimization The first attempt to solve TSP with neural networks is due to Hopfield and Tank [1985], which only scaled up to 30 cities. Recent developments of bespoke neural architectures [Vinyals et al., 2015, Vaswani et al., 2017] and performant hardware have made ML approaches increasingly efficient. Indeed, several architectures have been used to address CO problems, such as graph neural networks [Dai et al., 2017], recurrent neural networks [Nazari et al., 2018], and attention mechanisms [Deudon et al., 2018]. Kool et al. [2019] proposed an encoder-decoder architecture that we employ for TSP, CVRP and KP. The costly encoder is run once per problem instance, and the resulting embeddings are fed to a small decoder iteratively rolled out to get the whole trajectory, which enables efficient inference. This approach was furthered by Kwon et al. [2020] and Kim et al. [2022], who leveraged the underlying symmetries of typical CO problems (e.g. of starting positions and rotations) to realize improved training and inference performance using instance augmentations. Kim et al. [2021] also draw on Kool et al. and use a hierarchical strategy where a seeder proposes solution candidates, which are refined bit-by-bit by a reviser. Closer to our work, Xin et al. [2021] train multiple policies using a shared encoder and separate decoders.
Whilst this work (MDAM) shares our architecture and goal of training a population, our approach to enforcing diversity differs substantially. MDAM explicitly trades off performance with diversity by jointly optimizing policies and their KL divergence. Moreover, as computing the KL divergence for the whole trajectory is intractable, MDAM is restricted to only using it to drive diversity at the first timestep. In contrast, Poppy drives diversity by maximizing population-level performance (i.e. without any explicit diversity metric), uses the whole trajectory and scales better with the population size (we have used up to 32 agents instead of only 5).
Additionally, ML approaches usually rely on mechanisms to generate multiple candidate solutions [Mazyavkina et al., 2021]. One such mechanism consists in using improvement methods on an initial solution: de O. da Costa et al. [2020] use policy gradients to learn a policy that selects local operators (2-opt) given a current solution in TSP, while Lu et al. [2020] and Wu et al. [2021] extend this method to CVRP. This idea has been extended to enable searching a learned latent space of solutions [Hottung et al., 2021]. However, these approaches have two limitations: they are environment-specific, and the search procedure is inherently biased by the initial solution.
An alternative exploration mechanism is to generate a diverse set of trajectories by stochastically sampling a learned policy, potentially with additional beam search [Joshi et al., 2019], Monte Carlo tree search [Fu et al., 2021], dynamic programming [Kool et al., 2021], active search [Hottung et al., 2022], or simulation-guided search [Choo et al., 2022]. However, intuitively, the generated solutions tend to remain close to the underlying deterministic policy, implying that the benefits of additional sampled candidates diminish quickly.
Population-Based RL Populations have already been used in RL to learn diverse behaviors. In a different context, Gupta et al. [2018], Eysenbach et al. [2019], Hartikainen et al. [2020] and Pong et al. [2020] use a single policy conditioned on a set of goals as an implicit population for unsupervised skill discovery. Closer to our approach, another line of work revolves around explicitly storing a set of distinct policy parameters. Hong et al. [2018], Doan et al. [2020], Jung et al. [2020] and Parker-Holder et al. [2020] use a population to achieve a better coverage of the policy space. However, they enforce explicit attraction-repulsion mechanisms, which is a major difference with respect to our approach where diversity is a pure byproduct of performance optimization.
Our method is also related to approaches combining RL with evolutionary algorithms [EA; Khadka and Tumer, 2018, Khadka et al., 2019, Pourchot and Sigaud, 2019], which benefit from the sample-efficient RL policy updates while enjoying evolutionary population-level exploration. However, the population is a means to learn a unique strong policy, whereas Poppy learns a set of complementary strategies. More closely related, Quality-Diversity [QD; Pugh et al., 2016, Cully and Demiris, 2018] is a popular EA framework that maintains a portfolio of diverse policies. Pierrot et al. [2022] has recently combined RL with a QD algorithm, Map-Elites [Mouret and Clune, 2015]; unlike Poppy, QD methods rely on handcrafted behavioral markers, which is not easily amenable to the CO context.
One of the drawbacks of population-based RL is its expensive cost. However, recent approaches have shown that modern hardware, as well as targeted frameworks, enable efficient vectorized population training [Flajolet et al., 2022], opening the door to a wider range of applications.
3.1 Background and Motivation
RL Formulation A CO problem instance  $\rho$  sampled from some distribution  $\mathcal{D}$  consists of a discrete set of  $N$  variables (e.g. city locations in TSP). We model a CO problem as a Markov decision process (MDP) defined by a state space  $\mathcal{S}$ , an action space  $\mathcal{A}$ , a transition function  $T$ , and a reward function  $R$ . A state is a trajectory through the problem instance  $\tau_{t} = (x_{1},\ldots ,x_{t})\in \mathcal{S}$  where  $x_{i}\in \rho$ , and thus consists of an ordered list of variables (not necessarily of length  $N$ ). An action,  $a\in \mathcal{A}\subseteq \rho$ , consists of choosing the next variable to add; thus, given state  $\tau_{t} = (x_{1},\dots,x_{t})$  and action  $a$ , the next state is  $\tau_{t + 1} = T(\tau_t,a) = (x_1,\dots,x_t,a)$ . Let  $\mathcal{S}^*\subseteq \mathcal{S}$  be the set of solutions, i.e. states that comply with the problem's constraints (e.g., a sequence of cities such that each city is visited once and ends with the starting city in TSP). The reward function  $R:S^{*}\to \mathbb{R}$  maps solutions into scalars. We assume the reward is maximized by the optimal solution (e.g.  $R$  returns the negative tour length in TSP).
A policy  $\pi_{\theta}$  parameterized by  $\theta$  can be used to generate solutions for any instance  $\rho \sim \mathcal{D}$  by iteratively sampling the next action  $a \in \mathcal{A}$  according to the probability distribution  $\pi_{\theta}(\cdot \mid \rho, \tau_t)$ . We learn  $\pi_{\theta}$  using REINFORCE [Williams, 1992]. This method aims at maximizing the RL objective  $J(\theta) \doteq \mathbb{E}_{\rho \sim \mathcal{D}} \mathbb{E}_{\tau \sim \pi_{\theta}, \rho} R(\tau)$  by adjusting  $\theta$  such that good trajectories are more likely to be sampled in the future. Formally, the policy parameters  $\theta$  are updated by gradient ascent using  $\nabla_{\theta} J(\theta) = \mathbb{E}_{\rho \sim \mathcal{D}} \mathbb{E}_{\tau \sim \pi_{\theta}, \rho} (R(\tau) - b_{\rho}) \nabla_{\theta} \log(p_{\theta}(\tau))$ , where  $p_{\theta}(\tau) = \prod_{t} \pi_{\theta}(a_{t+1} \mid \rho, \tau_t)$  and  $b_{\rho}$  is a baseline. The gradient of the objective,  $\nabla_{\theta} J$ , can be estimated empirically using Monte Carlo simulations.
Figure 1: In this environment, the upward path always leads to a medium reward, while the left and right paths are intricate such that either one may lead to a low reward or high reward with equal probability. Left: An agent trained to maximize its expected reward converges to taking the safe upward road since acting optimally is too computationally intensive, as it requires solving the maze. Right: A 2-agent population can always take the left and right paths and thus get the largest reward.
Motivating Example We argue for the benefits of training a population using the example in Figure 1. In this environment, there are three actions: Left, Right, and Up. Up leads to a medium reward, while Left/Right lead to low/high or high/low rewards (the configuration is determined with equal probability at the start of each episode). Crucially, the left and right paths are intricate, so the agent cannot easily infer from its observation which one leads to a higher reward. Then, the best strategy for a computationally limited agent is to always go Up, as the guaranteed medium reward (2 scoops) is higher than the expected reward of guessing left or right (1.5 scoops). In contrast, two agents in a population can go in opposite directions and always find the maximum reward. There are two striking observations: (i) the agents do not need to perform optimally for the population performance to be optimal (one agent gets the maximum reward), and (ii) the performance of each individual agent is worse than in the single-agent case.
The discussed phenomenon can occur when (i) some optimal actions are too difficult to infer from observations and (ii) choices are irreversible (i.e. it is not possible to recover from a sub-optimal decision). This problem setting can be seen as a toy model for the challenge of solving an NP-hard CO problem in a sequential decision-making setting. In the case of TSP, for example, the number of possible unique tours that could follow from each action is exponentially large and, for any reasonably finite agent capacity, essentially provides the same obfuscation over the final returns. In this situation, as shown above, maximizing the performance of a population will require agents to specialize and likely yield better results than in the single-agent case.
3.2 Poppy
We present the components of Poppy: an RL objective encouraging agent specialization, and an efficient training procedure taking advantage of a pre-trained policy.
Population-Based Objective At inference, reinforcement learning methods usually sample several candidates to find better solutions. This process, though, is not anticipated during training, which optimizes the 1-shot performance with the usual RL objective  $J(\theta) = \mathbb{E}_{\rho \sim \mathcal{D}}\mathbb{E}_{\tau \sim \pi_{\theta},\rho}R(\tau)$  previously presented in Section 3.1. Intuitively, given  $K$  trials, we would like to find the best set of policies  $\{\pi_1,\dots ,\pi_K\}$  to rollout once on a given problem. This gives the following population objective:
$$
J _ {\text {p o p}} \left(\theta_ {1}, \dots , \theta_ {K}\right) \doteq \mathbb {E} _ {\rho \sim \mathcal {D}} \mathbb {E} _ {\tau_ {1} \sim \pi_ {\theta_ {1}}, \dots , \tau_ {K} \sim \pi_ {\theta_ {K}}} \max  \left[ R \left(\tau_ {1}\right), \dots , R \left(\tau_ {K}\right) \right],
$$
where each trajectory  $\tau_{i}$  is sampled according to the policy  $\pi_{\theta_i}$ . Maximizing  $J_{\mathrm{pop}}$  leads to finding the best set of  $K$  agents which can be rolled out in parallel for any problem.
Theorem 1 (Policy gradient for populations). The gradient of the population objective is:
$$
\nabla J _ {\mathrm {p o p}} (\theta_ {1}, \theta_ {2}, \ldots , \theta_ {K}) = \mathbb {E} _ {\rho \sim \mathcal {D}} \mathbb {E} _ {\tau_ {1} \sim \pi_ {\theta_ {1}}, \ldots , \tau_ {K} \sim \pi_ {\theta_ {K}}} \big (R (\tau_ {i ^ {*}}) - R (\tau_ {i ^ {* *}}) \big) \nabla \log p _ {\theta_ {i ^ {*}}} (\tau_ {i ^ {*}}),
$$
where:  $i^{*} = \arg \max_{i\in \{1,\dots ,K\}}\left[R(\tau_{i})\right]$  (index of the agent that got the highest reward) and  $i^{**} = \arg \max_{i\neq i^{*}}\left[R(\tau_{i})\right]$  (index of the agent that got the second highest reward).
The proof is provided in Appendix B.1. Remarkably, it corresponds to rolling out every agent and only training the one that got the highest reward on any problem. This formulation applies across various problems and directly optimizes for population-level performance without explicit supervision or handcrafted behavioral markers.
The theorem trivially holds when instead of a population, one single agent is sampled  $K$  times, which falls back to the specific case where  $\theta_{1} = \theta_{2} = \dots = \theta_{K}$ . In this case, the gradient becomes:
$$
\nabla J _ {\mathrm {p o p}} (\theta) = \mathbb {E} _ {\rho \sim \mathcal {D}} \mathbb {E} _ {\tau_ {1} \sim \pi_ {\theta}, \dots , \tau_ {K} \sim \pi_ {\theta}} \left(R \left(\tau_ {i ^ {*}}\right) - R \left(\tau_ {i ^ {* *}}\right)\right) \nabla \log p _ {\theta} \left(\tau_ {i ^ {*}}\right),
$$
Remark. The occurrence of  $R(\tau_{i^{**}})$  in the new gradient formulation can be surprising. However, Theorem 1 can be intuitively understood as training each agent on its true contribution to the population's performance: if for any problem instance the best agent  $i^*$  was removed, the population performance would fall back to the level of the second-best agent  $i^{**}$ , hence its contribution is indeed  $R(\tau_{i^{*}}) - R(\tau_{i^{**}})$ .
Optimizing the presented objective does not provide any strict diversity guarantee. However, note that diversity maximizes our objective in the highly probable case that, within the bounds of finite capacity and training, a single agent does not perform optimally on all subsets of the training distribution. Therefore, intuitively and later shown empirically, diversity emerges over training in the pursuit of maximizing the objective.
Algorithm 1 Poppy training
1: Input: problem distribution  $\mathcal{D}$ , number of agents  $K$ , batch size  $B$ , number of training steps  $H$ , pre-trained parameters  $\theta$ .
2:  $\theta_{1},\theta_{2},\ldots ,\theta_{K}\gets \mathsf{CLONE}(\theta)$  {Clone the pre-trained agent  $K$  times.}
3: for step 1 to  $H$  do
4:  $\rho_{i}\gets \mathsf{Sample}(\mathcal{D})\forall i\in 1,\ldots ,B$
5:  $\tau_{i}^{k}\gets \mathrm{Rollout}(\rho_{i},\theta_{k})\forall i\in 1,\ldots ,B,\forall k\in 1,\ldots ,K$
6:  $k_{i}^{*}\gets \arg \max_{k\leq K}R(\tau_{i}^{k})\forall i\in 1,\ldots ,B$  {Select the best agent for each problem  $\rho_{i}$  -}
7:  $\nabla L(\theta_1,\dots ,\theta_K)\gets \frac{1}{B}\sum_{i <   B}\mathrm{REINFORCE}(\tau_i^{k_i^*})$  {Backpropagate through these only.}
8:  $(\theta_{1},\ldots ,\theta_{K})\gets (\theta_{1},\ldots ,\theta_{K}^{-}) - \alpha \nabla L(\theta_{1},\ldots ,\theta_{K})$
Training Procedure The training procedure consists of two phases:
1. We train (or reuse) a single agent using an architecture suitable for solving the CO problem at hand. We later outline the architectures used for the different problems.
2. The agent trained in Phase 1 is cloned  $K$  times to form a  $K$ -agent population. The population is trained as described in Algorithm 1: only the best agent is trained on any problem. Agents implicitly specialize in different types of problem instances during this phase.
Phase 1 enables training the model without the computational overhead of a population. Moreover, we informally note that applying the Poppy objective directly to a population of untrained agents can be unstable. Randomly initialized agents are often ill-distributed, hence a single (or few) agent(s) dominate the performance across all instances. In this case, only the initially dominating agents receive a training signal, further widening the performance gap. Whilst directly training a population
Figure 2: Phases of the training process with a model using static instance embeddings. Left (Phase 1): the encoder and the decoder are trained from scratch. Right (Phase 2): the decoder is cloned  $K$  times, and the whole model is trained using the Poppy training objective (i.e. the gradient is only propagated through the decoder that yields the highest reward).
of untrained agents for population-level performance may be achievable with suitable modifications, we instead opt for the described pre-training approach as it is efficient and stable. Throughout this article, we use the reinforcement learning baseline of POMO [Kwon et al., 2020], which was proved to be efficient when training with multiple starting points. We investigate in Appendix C the effect of using the analytical baseline from Theorem 1 instead, which we show leads to improved performance.
Architecture To reduce the memory footprint of the population, some of the model parameters can be shared. For example, the architecture for TSP, CVRP, and KP uses the attention model by Kool et al. [2019], which decomposes the policy model into two parts: (i) a large encoder  $h_{\psi}$  that takes an instance  $\rho$  as input and outputs embeddings  $\omega$  for each of the variables in  $\rho$ , and (ii) a smaller decoder  $q_{\phi}$  that takes the embeddings  $\omega$  and a trajectory  $\tau_t$  as input and outputs the probabilities of each possible action. Figure 2 illustrates the training phases of such a model. For JSSP, we implement a similar encoder-decoder architecture taken from Jumanji [Bonnet et al., 2023] with the encoder being shared across the population. For all problems, we build a population of size  $K$  by sharing the encoder  $h_{\psi}$  with all agents, and implementing agent-specific decoders  $q_{\phi_i}$  for  $i \in \{1, \dots, K\}$ . This is motivated by (i) the encoder learning general representations that may be useful to all agents, and (ii) reducing the parameter overhead of training a population by keeping the total number of parameters close to the single-agent case. Please see Appendix A.1 for a discussion on model sizes.
We evaluate Poppy on four CO problems: TSP, CVRP, KP and JSSP. We use the JAX implementations from Jumanji [Bonnet et al., 2023] to leverage hardware accelerators (e.g. TPU). To emphasize the generality of our method, we use the hyperparameters from Kwon et al. [2020] for each problem when possible. We run Poppy for various population sizes to demonstrate its time-performance tradeoffs.