POMO: Policy Optimization with Multiple Optima for Reinforcement Learning

POMO begins with designating $N$ different nodes $\{a_{1}^{1},a_{1}^{2},\dots,a_{1}^{N}\}$ as starting points for exploration (Figure 2, (b)). The network samples $N$ solution trajectories $\{\bm{\uptau}^{1},\bm{\uptau}^{2},\dots,\bm{\uptau}^{N}\}$ via Monte-Carlo method, where each trajectory is defined as a sequence

In previous work that uses RNN- or Transformer-style neural architectures, the first node from multiple sample trajectories is always chosen by the network. A trainable START token, a legacy from NLP where those models originate, is fed into the network, and the first node is returned (Figure 2, (a)). Normally, the use of such START token is sensible, because it allows the machine to learn to find the “correct” first move that leads to the best solution. In the presence of multiple “correct” first moves, however, it forces the machine to favor particular starting points, which may lead to a biased strategy.

When all first moves are equally good, therefore, it is wise to apply entropy maximization techniques [27] to improve exploration. Entropy maximization is typically carried out by adding an entropy regularization term to the objective function of RL. POMO, however, directly maximizes the entropy on the first action by forcing the network to always produce multiple trajectories, all of them contributing equally during training.

Note that these trajectories are fundamentally different from repeatedly sampled $N$ trajectories under the START token scheme [28]. Each trajectory originating from a START token stays close to a single optimal path, but $N$ solution trajectories of POMO will closely match $N$ different node-sequence representations of the optimal solution. Conceptually, explorations by POMO are analogous to guiding a student to solve the same problem repeatedly from many different angles, exposing her to a variety of problem-solving techniques that would otherwise be unused.

POMO is based on the REINFORCE algorithm [20]. Once we sample a set of solution trajectories $\{\bm{\uptau}^{1},\bm{\uptau}^{2},\dots,\bm{\uptau}^{N}\}$, we can calculate the return (or total reward) $R(\bm{\uptau}^{i})$ of each solution $\bm{\uptau}^{i}$. To maximize the expected return $J$, we use gradient ascent with an approximation

where $p_{\theta}(\bm{\uptau}^{i}|s)\equiv\prod_{t=2}^{M}p_{\theta}(a_{t}^{i}|s,a^{i}_{1:t-1})$.

In Equation (3), $b^{i}(s)$ is a baseline that one has some freedom of choice to reduce the variance of the sampled gradients. In principle, it can be a function of $a_{1}^{i}$, assigned differently for each trajectory $\bm{\uptau}^{i}$. In POMO, however, we use the shared baseline,

Algorithm 1 presents the POMO training with mini-batch.

POMO baseline induces less variance in the policy gradients compared to the greedy-rollout baseline [10]. The advantage term in Equation (3), $R(\bm{\uptau}^{i})-b^{i}(s)$, has zero mean for POMO, whereas the greedy-rollout baseline results in negative advantages most of the time. This is because sample-rollouts (following softmax of the policy) have difficulty in surpassing greedy-rollouts (following argmax of the policy) in terms of the solution qualities, as will be demonstrated later in this paper. Also, as an added benefit, POMO baseline can be computed efficiently compared to other baselines used in previous deep-RL construction methods, which require forward passes through either a separately trained network (Critic [7, 8]) or the cloned policy network (greedy-rollout [10]).

Most importantly, however, the shared baseline used by POMO makes RL training highly resistant to local minima. After generating $N$ solution trajectories $\{\bm{\uptau}^{1},\bm{\uptau}^{2},\dots,\bm{\uptau}^{N}\}$, if we do not use the shared baseline but strictly stick to the greedy-rollout baseline scheme [10] instead, each sample-rollout $\bm{\uptau}^{i}$ would be assessed independently. Actions that produced $\bm{\uptau}^{i}$ would be reinforced solely by how much better (or worse) it performed compared to its greedy-rollout counterpart with the same starting node $a_{1}^{i}$. Because this training method is guided by the difference between the two rollouts produced by two closely-related networks, it is likely to converge prematurely at a state where both the actor and the critic underperform in a similar fashion. With the shared baseline, however, each trajectory now competes with $N-1$ other trajectories where no two trajectories can be identical. With the increased number of heterogeneous trajectories all contributing to setting the baseline at the right level, premature converge to a suboptimal policy is heavily discouraged.

Construction type neural net models for CO problems have two modes for inference in general. In “greedy mode,” a single deterministic trajectory is drawn using argmax on the policy. In “sampling mode,” multiple trajectories are sampled from the network following the probabilistic policy. Sampled solutions may return smaller rewards than the greedy one on average, but sampling can be repeated as many times as needed at the computational cost. With a large number of sampled solutions, some solutions with rewards greater than that of the greedy rollout can be found.

Using the multi-starting-node approach of POMO, however, one can produce not just one but multiple greedy trajectories. Starting from $N$ different nodes $\{a_{1}^{1},a_{1}^{2},\dots,a_{1}^{N}\}$, $N$ different greedy trajectories can be acquired deterministically, from which one can choose the best similarly to the “sampling mode” approach. $N$ greedy trajectories are in most cases superior than $N$ sampled trajectories.

We implement POMO by assigning every node to be a starting point for a sample rollout for TSP. That is, the number of starting nodes ($N$) we use is 20 for TSP20, 50 for TSP50, and 100 for TSP100.

In Table 2 we compare the performance of POMO on TSP with other baselines. The first group of baselines shown at the top are results from Concorde [32] and a few other representative non-learning-based heuristics. We have run Concorde ourselves to get the optimal solutions, and other solvers’ data are adopted from Wu et al. [11] and Kool et al. [10]. The second group of baselines are from deep RL improvement-type approaches found in the literature [9, 11, 12]. In the third group, we present the results from the AM that is trained by our implementation of REINFORCE with a greedy rollout baseline [10] instead of POMO.

Given 10,000 random instances of TSP20 and TSP50, POMO finds near-optimal solutions with optimality gaps of 0.0006% in seconds and 0.025% in tens of seconds, respectively. For TSP100, POMO achieves the optimality gap of 0.14% in a minute, outperforming all other learning-based heuristics significantly, both in terms of the quality of the solutions and the time it takes to solve.

In the table, results under “AM, greedy” method and “POMO, single trajec.” method are both from the identical network structure that is tested by the same inference technique. The only difference was training, so the substantial improvement (e.g. from 3.51% to 1.07% in optimality gap on TSP100) indicates superiority of the POMO training method. As for the inference techniques, it is shown that the combined use of multiple greedy rollouts of POMO and the $\times 8$ instance augmentation can reduce the optimality gap even further, by an order of magnitude.

Learning curves of TSP50 and TSP100 in Figure 3 show that POMO training is more stable and sample-efficient. In reading these graphs, one should keep in mind that POMO uses $N$-times more trajectories than simple REINFORCE for each training epoch. POMO training time is, however, comparable to that of REINFORCE, thanks to the parallel processing on trajectory generation. For example, TSP100 training takes about 7 minutes per epoch for POMO while it take 6 minutes for REINFORCE.

When POMO trains a policy net, ideally it should use only the “good” starting nodes from which one can roll out optimal solutions. But, unlike TSP, not all nodes in CVRP can be the first steps for optimal trajectories (see Figure 3) and there is no way of figuring out which of the nodes are good without actually knowing the optimal solution a priori. One way to resolve this issue is to add a secondary network that returns candidates for optimal starting nodes to be used by POMO. We leave this approach for future research, however, and in our CVRP experiment we stick to the same policy net that we have used for TSP without an upgrade. We simply use all nodes as starting nodes for POMO exploration regardless of whether they are good or bad.

This naive way of applying POMO can still make a powerful solver. Experiment results on CVRP with 20, 50, and 100 customer nodes are reported in Table 3, and POMO is shown to outperform simple REINFORCE by a large margin. Note that there is no algorithm yet that can find optimal solutions of 10,000 random CVRP instances in a reasonable time, so the “Gap” values in the table are given relative to LKH3 [33] results. POMO has a smaller gap in CVRP100 (0.32%) than in CVRP50 (0.45%), which is probably due to LKH3 falling faster in performance than POMO as the size of the problem grows.

Improvement-type neural approaches, such as L2I by Lu et al. [15], can produce better results than (single-run) LKH3, given long enough search time.⁴⁴4In previous versions of this paper, the runtimes for L2I in Table 3 were incorrect, and we have now fixed them. The runtimes we mistakenly quoted were for single-instances, rather than for 10k instances. To emphasize the differences between POMO and L2I (other than the speed), POMO is a general RL tool that can be applied to many different CO problems in a purely data-driven way. One the other hand, L2I is a specialized routing problem solver based on a handcrafted pool of improvement operators. Because POMO is a construction method, it is possible to combine it with other improvement methods to produce even better results.

We choose KP to demonstrate flexibility of POMO beyond routing problems. Similarly to the case of CVRP, we reuse the neural net for TSP and take the naive approach that uses all items given in the instance as the first steps for rollouts, avoiding an additional, more sophisticated “SelectStartNodes” network to be devised. In solving KP, a weight and a value of each item replaces x- and y-coordinate of each node of TSP. As the network generates a sequence of items, we put these items into the knapsack one by one until the bag is full, at which point we terminate the sequence generation.

In Table 4, the POMO results are compared with the optimal solutions based on dynamic programming, as well as those by greedy heuristics and our implementation of PtrNet [7] and the original AM method [10]. Even without the instance augmentation, POMO greatly improves the quality of the solutions one can acquire from a deep neural net.