Heterogeneous Graph Transformers for Simultaneous Mobile Multi-Robot Task Allocation and Scheduling under Temporal Constraints

We thank the reviewer for their detailed feedback and constructive comments. We have made revisions based on the suggestions.

There are some expression and grammar mistakes in the paper (e.g., line 51, 163 and 181), which should be avoided in academic writing.

We have corrected the grammar mistakes mentioned and further improved our manuscript.

Overall, the performance of TARGETNET is not outstanding enough compared to the baselines. It exhibits clear superiority in the computation time, but its optimal and successful rates are inferior to MILP solver in small, medium and large scale tasks. As MILP solver is supported by solid mathematical theory, users would prefer MILP in most cases, making the application of TARGETNET limited in extremely large cases.

We would like to thank the reviewer for their thorough analysis of the limitations of the Machine Learning-based approaches to optimization problems.

Our primary contribution is not to replace MILP in scenarios where it is feasible, but to provide a scalable solution that delivers high-quality, near-optimal schedules in a fraction of the time for problem sizes where MILP fails entirely. For instance, our experiments show that TARGETNET finds solutions orders of magnitude faster, making it practical for real-time applications where a good solution quickly is more valuable than a perfect solution that arrives too late.

Developing fast, high-quality approximate solvers like TARGETNET is a vital research direction that can work together with exact solvers. As noted in [1, 2], such methods can provide a high-quality "warm start" to dramatically reduce the search space and runtime of a MILP solver. Therefore, we believe TARGETNET is not "limited" but rather is a significant and practical step forward for a broad and critical class of computationally challenging optimization problems.

We have also conducted additional experiments, including TARGETNET without edge attention, TARGETNET\E, and a modification of TARGETNET with a non-normalized edge attention. The results, presented as part of the rebuttal for Reviewer EqLo's comments, show that one of the three seeds is able to converge into a policy that is within 5% of the optimal fully feasible solution on an extra-large scale. The modification of TARGETNETv2 allows a more robust model that can, on average, converge to a policy that is better than other learning based baselines, indicating a less brittle model that depends on the initial model weights. These results will be added to the camera-ready version.

[1] Bengio, Y., Lodi, A., & Prouvost, A. (2021). Machine learning for combinatorial optimization: a methodological tour d’horizon. European Journal of Operational Research, 290(2), 405-421. [2] Kruber, M., Lübbecke, M. E., & Parmentier, A. (2017, May). Learning when to use a decomposition. In International conference on AI and OR techniques in constraint programming for combinatorial optimization problems (pp. 202-210). Cham: Springer International Publishing.

Besides the agents and tasks as nodes, the authors further include |A||T| assignment nodes in the graph. When encountered with extremely large scale tasks, the complexity of the graph would increase significantly. However, the experiment results show that the computation time of TARGETNET is lower than baselines. Could the authors explain why TARGETNET maintains a high efficiency?

TARGETNET only constructs a single graph representation of the world, including the heterogenous travel time from one task to another task for each agent, which is represented as an edge between the Assignment Nodes, with a complexity of $O(|A||T|^2)$ number of edges due to the limitation of per agent travel.

Unlike baselines that still require the simulation of the agent movement to new location to update agent makespan and when they would be available to move to the next task location, the simultaneous nature of TARGETNET allow for the generation of the entire schedule without needing to interact with the environment, relying on the initial state observation, $s_0$, to generate a full schedule.

Furthermore, TARGETNET scales polynomially with the number of agents and tasks while the problem itself is NP-hard. This is one of the key advantages of learning-based models, as they can learn potential heuristics that outperform existing methods.

As shown in Figure 3 and 4, the performance of TARGETNET is severely declined after removing the residual connections. Is this indicating that the improved performance of TARGETNET is due to the residual connections rather than the graph-based learning methodology?

We would like to refer the Reviewer to the rebuttal to the feedback of Reviewer EqLo for additional experiments for our model without edge attention, TARGETNET\E. We agree with the reviewer that the performance of the Graph Models has a significant increase in the performance with the addition of Residual Connections, and the presence of the Edge Features makes the model less brittle, capable of learning a better mean performance across 3 different seeds (10, 11, 12).

Furthermore, our additional experiments with TARGETNETv2 that allow a non-normalized edge attention encapsulate the policies that can be learned by TARGETNET\E, as they are able to learn the same policies. This makes the addition of Edge Attention lead to a more robust model that can learn a wide range of policies that models without edge attention are not able to represent.

We are currently running the training and evaluation of the Sequential variant of the model without edge features, Seq-TARGETNET\E. Due to the slower nature of the model, we will provide these results in the camera-ready version of our work along with a detailed discussion.

We sincerely thank the reviewer for their detailed feedback and insightful comments. We have incorporated revisions based on this feedback.

The proposed method has a fast inference speed, while the training time is not mentioned. Can you discuss the time cost of training stage of TARGETNET? The training time was not mentioned as each model is trained once in the small scale and tested on a range of different scales.

We have conducted an additional evaluation on our system for learning-based baselines for over 100 epochs. We report the mean and standard deviation of time in seconds to train each model for a single epoch below:

As seen in the training time, the training speed of different models is consistent with the inference speed. Simultaneous decision-making models of TARGETNET are faster by a factor of x20. As the models are trained on 20 task problems, the empirical results align with the theoretical complexity of $O(|A||T|^2)$ for sequential models and $O(|A||T|) for simultaneous models. The increase in the feature space size with the presence of residuals and edge attention leads to an increase in training and evaluation time.

The training time presented above also implies that it is faster to train a simultaneous TARGETNET over the existing sequential models from prior work, leading to training and deploying models within the time frame that it takes an exact solver to produce a single suboptimal solution for the extra-large scale.

The interpretability of the proposed method is limited, which is a common problem in deep RL.

We agree with the Reviewer that Interpretability is a common problem in both Deep RL and Optimization Problems. We further discuss it in paragraph starting in Line 329 in Section 6.

a. Using RL to solve combinatorial optimization problems (e.g. graph-based problems) is not a new idea [1,2,3].

We agree that there are prior works that focus on using RL to solve combinatorial optimization problems. The method presented in TARGETNET, which uses RL for both Task Allocation and Scheduling in a single forward propagation, is novel.

Please provide more discussion on the difference in problems and methods between this paper and [4,5,6].

We would like to thank the reviewer for referring the related work which we will discuss more in depth. We agree that contextualizing our work within the broader field of RL for Combinatorial Optimization (CO) is crucial. We have revised our related work section to better position TARGETNET and now include a detailed discussion of the suggested papers. The key distinctions are:

• [4] Ma et al., 2019: This work proposes a hierarchical GNN for the Traveling Salesman Problem with Time Windows (TSPTW). While related, our HM-MATAS formulation addresses a more complex problem, equivalent to a multi-agent TSP variant (m-TSPTW) with additional heterogeneity constraints for agents and tasks. TARGETNET is designed to handle the joint allocation and scheduling aspects of this multi-agent, multi-task setting in a single pass, and accounts for heterogeneous task durations along with the heterogeneous travel time.

• [5] Altundas et al. & [6] Wang and Gombolay: The problem formulation of both [5] and [6] assumes that travel times are negligible compared to task service durations, leading to a simpler graph representation of the problem. Our framework, TARGETNET, is explicitly designed to handle scenarios where heterogeneous travel times are a critical component of the scheduling decision and may be a key factor in success. Furthermore, [5] relies on an ensemble of policies to achieve its best performance, whereas TARGETNET generates a high-quality schedule in a single forward pass. [6] focuses on a stochastic maintenance scheduling problem, leveraging the HetGAT model we benchmark against. However, the key challenge in [6] is when to dispatch agents to tasks that appear over time, which differs from our deterministic generation of complete schedules in fully observable settings.

We have incorporated these comparisons into our revised manuscript.

It is observed that the sequential decision making version of TARGETNET have better performance than the one-step version. Some prior work [7] has shown the potential to boost the performance of stateless policy by a simple distillation process, which means we can distill the sequential policy to a one-step policy by a simple behavior cloning process. Can you try this simple distillation process to boost the performance of the one-step version of TARGETNET?

The proposed distillation is indeed a powerful technique for creating a fast and high-performance policy. However, the primary motivation for our one-step model is to demonstrate its capability to learn through Reinforcement Learning. The goal is to provide a lightweight alternative for scenarios with limited computational budgets. The distillation process, while effective, would first require training the full sequential model, thereby negating the training-time advantages that make the one-step model appealing.

Our paper's focus is on analyzing the trade-off between the computationally expensive sequential approach and the efficient one-step approach. While distilling the sequential policy is an excellent direction for future work to create a fast inference model, it addresses a different research question than the one we focus on.

Furthermore, the additional experiments conducted, as presented in the rebuttal for Reviewer EqLo's comments, show that simultaneous models are able to learn policies that outperform sequential models in every scale while being more robust to initial starting weights.

The edge attention mechanism to incorporate the edge information seems to be an important contribution of this paper. However, there is a simple way to incorporate the edge information into the node attention mechanism: we can create a new node for each edge with the edge information as the node feature. This way, we can use the same node attention mechanism to incorporate the edge information. Can you discuss the difference between this simple way and the proposed edge attention mechanism? And what if the baseline method Res-HetGAT uses this simple way to incorporate the edge information?

We thank the reviewer for their suggested feedback. While the edge features can be represented as a node feature with the addition of new nodes for each node, the added node would lead to the need for additional layers for message passing. Furthermore, the two-hop message passing for the node features would lead to added complexities and potential for information loss. Our proposed method also combines the edge and node features when updating the node features of the next layer, leveraging the linearity of attention such that $A(B || C) = A_1 B + A_2 C$, where $A = (A_1||A_2)$, for attention vectors, $A$, $A_1$ and $A_2$ and feature vectors $B$ and $C$.

For both HetGAT and HGT, the simple way of incorporating the edge features and attention would lead to increased complexity and need for additional layers, leading to increased computation time that may surpass the method proposed in our work.

Based on the feedback of the reviewer, we have also improved our model, adding TARGETNETv2 into our results. TARGETNETv2 removes the softmax of attention heads in the edge feature, allowing for our model to represent models with and without edge attention by learning the attention value. This leads to a more versatile and robust model in the presence of relational information, such as heterogeneous travel times. Again, we would like to refer the reviewer to the rebuttal for Reviewer EqLo's comments for these additional results.

[1] Wagner, Adam Zsolt. “Constructions in combinatorics via neural networks.” ArXiv abs/2104.14516, 2021.

[2] Darvariu V A, Hailes S, Musolesi M. “Graph reinforcement learning for combinatorial optimization: A survey and unifying perspective” [J]. arXiv preprint arXiv:2404.06492, 2024.

[3] Peng Y, Choi B, Xu J. “Graph learning for combinatorial optimization: a survey of state-of-the-art”[J]. Data Science and Engineering, 2021.

[4] Ma, Qiang, Suwen Ge, Danyang He, Darshan D. Thaker and Iddo Drori. “Combinatorial Optimization by Graph Pointer Networks and Hierarchical Reinforcement Learning.” ArXiv abs/1911.04936, 2019.

[5] Batuhan Altundas, Zheyuan Wang, Joshua Bishop, and Matthew Gombolay. “Learning Coordination Policies over Heterogeneous Graphs for Human-Robot Teams via Recurrent Neural Schedule Propagation.”

[6] Zheyuan Wang and Matthew Gombolay. “Stochastic Resource Optimization over Heterogeneous Graph Neural Networks for Failure-Predictive Maintenance Scheduling”.

We thank the reviewer for their detailed feedback and constructive suggestions. Below, we address the noted weaknesses and questions.

Lack of interpretability — like most deep GNNs, the learned policy is hard to analyze or debug in safety-critical applications.

We acknowledge that interpretability is a limitation of our model, and discuss it in the limitations section. Future Work discusses potential approaches to improve the interpretability while leveraging the advantages of the model.

Slightly weaker performance than sequential models in extremely large-scale cases, which can benefit from per-step re-optimization (but at a huge cost in time).

In Section 6, we propose potential future work in utilizing a meta-policy that combines the performance of the sequential models with simultaneous models. Furthermore, we have conducted additional experiments, leading to increase in performance of our model TARGETNETv2 to near optimal levels while being less brittle than the models not using edge attention. These results can be seen below.

The design of transformer seems very specific to the nature of MATAS. I wish to see a more principled design like incorporating next-token prediction, and test the scaling law, etc.

The suggestion of next-token prediction is rooted in autoregressive models like LLMs, which excel at sequential data. However, the Multi-Agent Task Allocation and Scheduling (MATAS) problem is fundamentally a combinatorial optimization problem on a graph, not a sequence generation task. While Transformers can be used in the context of combinatorial optimization problems [1], the approach of graph attention are applied to node and edge features that are connected to individual nodes on a graph, allow sparse convolution that can scale to larger problems [2].

We would like to refer the reviewer to Hu et al. 2020 [3] for full explanation of the Heterogenous Graph Transformers (HGT) as well as Appendix 2. While the HGT leverage the attention mechanism popularized by Vaswani et al. 2017 [4], our model does not utilize token prediction in the same way that a transformer architecture would achieve.

HGT is a common Graph Neural Network algorithm that leverages a distinct attention mechanism. To our knowledge, the edge attention mechanism we propose in our paper has not been integrated into the HGT architecture. While we leverage the method for the MATAS problems, further research would need to be conducted for any other problem where the relational edge features are encoded in the input of a model.

We train our model in small scale problems and evaluate the trained models in large scale problem without any fine tuning, as such the presented results do not have a change in the parameter counts during training. Future work will involve training in different scale environments to test how the model training scales, however this is beyond the scope of this paper.

There are ablation studies on TARGETNET\ER and TARGETNET\R, how about TARGETNET\E?

Below is the performance of the best of three seeds. Along with TARGETNET\E, we introduce TARGETNETv2, which is a modification of TARGETNET where edge attention is not normalized using the softmax function. The non-normalized attention allows for TARGETNETv2 to learn policies that can be learned by TARGETNET\E, by learning the edge attention weight to 0. The modification allows our model to learn:

Below is the mean and standard deviation of the performance of the learning-based models using three seeds (10, 11, 12) as an expansion of the ablation study presented in our paper.

While the TARGETNET\E is able to learn a near-optimal policy in one instance of the seed, leading to high performance when the best performing schedule is chosen based on models trained on different seeds, the mean and standard deviation of the performance across seeds show that TARGETNETv2, which removes the softmax of the attention, can not only learn the same policies, but also converge on other locally optimal scheduling policies, making TARGETNETv2 less brittle.

Why does sequential version performs better on the largest task? What are the typical failure modes of TARGETNET under this setting?

We discuss the advantages of sequential version in Line 290. At the extra large scale, the action space for Task Allocation, which is $O(|A||T|)$, becomes as large as 16000, which leads to the probability distribution to get close to uniform, leading to a drop in performance.

[1] Z. Yang, A. Ishay, and J. Lee, “Learning to Solve Constraint Satisfaction Problems with Recurrent Transformer,” Jul. 10, 2023, arXiv:2307.04895. arXiv.2307.04895

[2] Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., & Bengio, Y. (2017). “Graph attention networks”. arXiv:1710.10903

[3] Hu, Z., Dong, Y., Wang, K., & Sun, Y. (2020, April). “Heterogeneous graph transformer”. In Proceedings of the web conference 2020 (pp. 2704-2710).

[4] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). “Attention is all you need.” Advances in neural information processing systems.

We thank the reviewer for their positive assessment and constructive feedback.

The evaluation is on fairly contrived domains, which seem somewhat removed from real-world settings; related, the authors acknowledge their approach is not directly applicable to more dynamic settings, since the plans are pre-computed

Our primary motivation for designing this domain was to fill a specific gap in the literature for multi-scale, heterogeneous multi-agent scheduling in continuous spaces, which requires accounting for agent travel times, task durations, and complex inter-task dependencies. By synthesizing elements from well-established optimization problems like the Vehicle Routing Problem with Time Windows (VRPTW) and the Job Shop Scheduling Problem (JSSP), our domain, while simulated, is designed to capture the combinatorial complexity inherent in many real-world logistics and coordination challenges. We will clarify this connection in the paper to better contextualize our environment's relevance.

While the pre-computed plans limit applicability in highly dynamic environments, where frequent replanning is necessary, this assumption is standard for the class of offline task allocation and scheduling problems, including those addressed by our baselines, as well as more standard optimization problems such as VRP and Travelling Salesman Problem (TSP). The chosen baselines allows for a focused and fair evaluation of the scheduling algorithm itself, decoupled from the separate and significant challenges of real-time perception and collision avoidance. We believe our work provides a strong foundation for future extensions into more dynamic settings.

To make this clearer, we will expand the limitations section of the revised manuscript to better frame the problem setting and explicitly motivate dynamic re-planning as a key avenue for future work.

Nitpicks - Figure 2a text is fairly small; would recommend ensuring fontsize in diagrams are generally at the same scale of the caption.

We modified the Figure 2a to fix the fontsize and further ensured that the font size of the figures match the scale of the captions in the camera ready version.

Within the MARL literature, ALMA [1] is one method that may be relevant (it also involves allocating heterogenous agents to tasks dynamically; a learning-based approach); discussion on the differences in these settings, and the comparative merits of TARGETNET and ALMA would be useful. Can TARGETNET be applied in ALMA's setting as well? Or broadly, be integrated with more general, learning-based approaches? (as I reflect, it seems that ALMA's setting is more dynamic than the one investigated here)

ALMA[1] allows for allocation of agents to task in cooperative tasks followed by a subtask allocator to handle to order of assignments, using a Value-based model. In comparison, TARGETNET is a policy-based framework. While ALMA optimizes for performance of subtask for each subteam, TARGETNET optimizes for the global makespan and feasibility. Furthermore, TARGETNET focuses on generating scheduling prior to any action is taken based on the complete representation of the initial condition.

We will include ALMA[1] into the related section and discuss the differences in application and problem setup.

[1] Iqbal, S., Costales, R., & Sha, F. (2022). Alma: Hierarchical learning for composite multi-agent tasks. Advances in neural information processing systems, 35, 7155-7166.