Learning Memory-Enhanced Improvement Heuristics for Flexible Job Shop Scheduling

We appreciate the reviewer for their valuable feedback and address the main concerns below:

Analysis of embedding

In our framework, the policy network makes decisions based on both the current scheduling state (operation and machine embeddings) and historical information (historical action embedding). Specifically，

• Operation embeddings capture the current solution state. Since different scheduling solutions correspond to different graph structures[1]—i.e., different adjacency relations among operation nodes—we extract the topological structure of the graph and the contextual dependencies[1] among operations as the operation embeddings.
• Machine embeddings reflect the machine load by aggregating the features of neighboring operation nodes, and serve as critical heuristics for the policy to reassign operations from overloaded machines to less loaded ones[2], thereby reducing the makespan.
• Historical action embeddings represent actions taken under similar historical states. They serve as a reference to encourage diverse action selection and enhance exploration[3], when combined with the similarity penalty in the reward function.

Module selection in operation

Topological information and the semantic information of constraints are key to distinguishing different operation embeddings[1].

• For topological information, which distinguishes different schedule structures, we chose GIN[4]. Its proven strength, equivalent to the Weisfeiler-Lehman test, is essential for capturing the subtle but critical structural differences between operation nodes in varying solutions[1].
• For semantic information, we recognized that an operation's neighbors—job predecessors and machine predecessors—carry different semantic meanings[1, 2]. To handle this, we employed the GAT[6]. Its attention mechanism is perfectly suited to learn the varying semantic roles and importance of these different neighbor types, assigning weights accordingly and thus enriching the node representations.

Our module selection is not only theoretically motivated by the structural characteristics of FJSP graphs, but also aligned with mainstream practices[1, 2, 5, 7]. We conducted ablation studies by removing each component. Results below confirm that each of them is essential for learning informative and discriminative state representations that support effective policy learning.

While more advanced GNNs exist, our aim was to establish a strong, well-justified, and interpretable baseline. We believe this study provides a clear and robust validation of our architectural choices.

Dynamic change challenges

We are grateful for the reviewer's encouragement to discuss MIStar's broader applicability. Although our current work focuses on the standard FJSP, we designed MIStar with future dynamic adaptations in mind. We believe its core architectural strengths—a flexible graph representation, a size-agnostic GNN policy, and an adaptable reward function—make it well-suited to address the challenges mentioned.

1. Multi-capability Devices

The scenario where machines are capable of performing multiple types of actions naturally aligns with the modeling paradigm of FJSP. The various operations within a job typically correspond to different stages of processing, each representing a specific type of task[10]. In contrast to JSP, where each operation must be processed by a fixed machine[9](i.e., each machine performs only one type of action), FJSP allows each operation to be processed by multiple candidate machines[10]. This implies that machines can handle multiple processing stages within a job, thereby exhibiting multi-processing capabilities. Therefore, MIStar, as a method for solving FJSP, inherently supports the modeling and optimization of device functionality variability.

2. Machine quantity changes

Although changes in machine quantity may alter the graph structure, the size-agnostic nature of our GNN-based policy network enables MIStar to dynamically re-optimize[12] from any valid existing solution, without restarting the scheduling process from scratch. Furthermore, our memory mechanism remains robust, as the incidence matrices from schedules of varying sizes can be aligned via simple padding or cropping before similarity comparison. In addition, the reward function should not only focus on minimizing makespan but also incorporate schedule stability in re-scheduling scenarios. To this end, we can introduce a stability term[11] into the reward to minimize the differences in operation start times between the original and revised schedules $\sum_{j=1}^{|J|} \sum_{i=1}^{n_j} \left| st_t - st_{t-1} \right|$, thus encouraging minimal disruption to the former schedule[11,13].

To further handle different types of equipment changes, we introduce more specific strategies:

Machine Addition:

• New jobs can be integrated by expanding the state matrix and enriching feature vectors with arrival times. Their feature vectors should also be enriched with the arrival time.
• To avoid interference from irrelevant information, operations completed before the arrival of the new machine will not be recorded in the state. This helps reduce noise and ensures a compact state representation during re-scheduling with a large number of added machines.

Machine Reduction:

• Unavailable machines can be removed by pruning their corresponding entries in the state matrices. Alternatively, the original graph structure can be retained by setting the processing time on failed machines to infinity. In this case, the agent, which aims to minimize the makespan, will naturally avoid assigning operations to a machine marked as unavailable or having an infinite processing time.
• All operations currently assigned to failed machine must be re-assigned. Broken machines are masked out when generating candidate actions based on the Nopt2 neighborhood structure. As a result, invalid actions $[O_m, O_n, M_k]$ involving broken machines in the third dimension are effectively filtered out.

3. Shifts in priority

Job priorities may shift due to changes in customer demand, resulting in urgent jobs requiring higher priority[8]. This can be elegantly addressed through the reward function. Jobs can be grouped by urgency of customer orders—i.e., normal jobs and urgent jobs. The objective function for this scheduling problem can be augmented as the weighted sum of the makespan for normal jobs and the total tardiness for urgent jobs[8], which is to be minimized. Moreover, it is possible to divide jobs into more than two groups based on finer-grained urgency levels. The weights assigned to each group can be dynamically adjusted according to their priority. To adapt to real-time shifts in job priorities, the grouping and weighting can be periodically updated at fixed iteration intervals during the scheduling process.

On Open-Sourcing the code

We are fully committed to making our code and models publicly available upon acceptance of the paper to facilitate reproducibility and future research in the community.

We extend our heartfelt thanks to the reviewer for their detailed and thoughtful observations during the rebuttal process. Integrating these insights with the new experimental data and our discussions, we will enhance the final version of the paper to improve its clarity and presentation.

[1]Deep reinforcement learning guided improvement heuristic for job shop scheduling, 2022.
[2]Flexible job-shop scheduling via graph neural network and deep reinforcement learning, 2022.
[3]Marco: A memory-augmented reinforcement framework for combinatorial optimization, 2024.
[4]How powerful are graph neural networks?, 2018.
[5]A multi-action deep reinforcement learning framework for flexible Job-shop scheduling problem, 2022.
[6]Graph attention networks, 2017.
[7]Graph neural networks for job shop scheduling problems: A survey, 2025.
[8]Metaheuristics for a flow shop scheduling problem with urgent jobs and limited waiting times, 2021.
[9]Deterministic job-shop scheduling: Past, present and future, 1999.
[10]Routing and scheduling in a flexible job shop by tabu search, 1993.
[11]Dynamic scheduling in flexible job shop systems by considering simultaneously efficiency and stability, 2010.
[12]A review of dynamic scheduling: context, techniques and prospects, 2021.
[13]Multiobjective flexible job-shop rescheduling with new job insertion and machine preventive maintenance, 2022.

We thank the reviewer for their valuable suggestions and comments and address the main concerns below:

On the Memory Mechanism: Ablation Study and Design Justification

1. ablation study

We would like to clarify that an ablation study on the memory module was indeed conducted and is presented in Appendix I (Table 5). To directly answer the reviewer's concern, this study demonstrates that MIStar consistently outperforms its variant without the memory module, with performance gains ranging from 0.41% to 4.36% across different problem scales. This confirms the significant and positive contribution of our memory-enhanced design.

2. design of the state similarity metric

We appreciate the reviewer's sharp observation, which gives us an opportunity to clarify a subtle but crucial detail of our design. We commit to adding more details of $L_t$ to the revised version. The reviewer correctly notes that ignoring sequence information would be a major flaw. However, our operation-machine incidence matrix, $L_t$, does encode sequence information. Its elements are not merely binary indicators (0/1) of machine assignment, but rather represent the processing order index of each operation on its assigned machine.

This design ensures that two schedules with identical machine assignments but different sequences (e.g., a good sequence vs. a poor one) are correctly distinguished. For instance, a good sequence [1, 2, 3, 4, 5] and a poor, reversed sequence [5, 4, 3, 2, 1] on the same machine will be recognized as maximally dissimilar by the Frobenius inner product. The mathematical foundation for this lies in the Rearrangement Inequality, which guarantees that the inner product is maximized for similarly ordered sequences and minimized for oppositely ordered ones. This allows our similarity metric to be both computationally efficient and highly sensitive to the quality of the operation sequence.

Theorem: The Rearrangement Inequality states that, for two sequences $a_1 \le a_2 \le \dots \le a_n$ and $b_1 \le b_2 \le \dots \le b_n$, the inequalities

hold, where $\pi(1), \pi(2), \dots, \pi(n)$ is any permutation of $1, 2, \dots, n$.

3. alternative designs

We have considered and tested other expressive representations for the historical states that explicitly encode sequence information. For example:

1. The adjacency matrix $A_t^{M}$ of operations on machines clearly encodes the processing order between operations, while the binary operation-machine incidence matrix $L_t^{B}$ indicates machine assignment. The state similarity is then calculated by computing the inner products of both matrices separately and summing the results, expressed as:

2. For the operation-machine incidence matrices that store processing order indices, we compute the element-wise difference between two schedules, $L_t$ and $L_{t'}$, to obtain a difference matrix, and then use the reciprocal of its Frobenius norm to measure similarity, formulated as:

Dear reviewer,

As the discussion period is nearing its end, with fewer than three days left, I want to ensure we have addressed all your concerns to your satisfaction. If there are any additional points or feedback you'd like us to consider, please let us know at your earliest convenience. Your insights are invaluable to us, and we're eager to address any remaining issues promptly to further improve our work.

Thank you for your time and effort in reviewing our paper.

We appreciate the reviewers' feedback and address each comment below:

Dynamic events

We thank the reviewer for this important point. We would like to clarify that our focusing on static environments is not an oversight of dynamic problems, but aligns with the mainstream research paradigm in learning-based scheduling[1,2,3]. Even the static, deterministic FJSP remains a formidable challenge, and developing effective heuristics for it is an active and critical area of research[4].

That being said, we fully agree that adaptability to realistic dynamics is crucial. A key strength of MIStar is its modular architecture, designed for future extensibility and minimal modification to handle dynamic environments. Below, we detail how the graph representation, memory mechanism, and action space can be adapted for two representative examples.

Handling real-time job arrivals:

MIStar is exceptionally well-suited for dynamic re-optimization, a key advantage of our improvement-based approach over constructive methods.

1. Graph representation: The flexibility of our graph structure allows new jobs to be seamlessly integrated by expanding the state matrices to include their operations and associated constraints. To maintain focus and efficiency, completed operations can be pruned from the state, ensuring the agent always reasons over the most relevant part of the problem[3].
2. Memory mechanism: 0ur memory mechanism remains robust, as the incidence matrices from schedules of different sizes can be aligned via simple padding or cropping before similarity comparison.
3. Action space: Time-based constraints are easily handled. To respect arrival times, we simply filter out any action that would schedule a job's first operation before it is ready[5]. The problem's inherent precedence constraints, already encoded in our graph, automatically handle the sequencing for all subsequent operations.

Handling unexpected machine breakdowns:

1. Graph representation: Machine status is a feature in our graph representation, allowing the breakdown to be modeled by updating the node. The agent, trained to minimize makespan, will naturally avoid assigning operations to unavailable machines or those with infinite processing time. Alternatively, Unavailable machines can be directly removed from the graph by pruning their corresponding entries in the state matrices.
2. Memory mechanism: Machine breakdowns can be reflected by assigning a special value (e.g., –1) to unavailable columns in incidence matrices. Alternatively, memory entries can include labels (e.g., breakdown presence or broken machine count) for coarse-grained filtering, followed by fine-grained similarity matching.
3. Action space: Operations on failed machines must be reassigned. Unavailable machines are masked out in the Nopt2 neighborhood to filter invalid actions $[O_m, O_n, M_k]$ involving broken machines.

Finally, the size-agnostic nature of our GNN-based policy network[7] is a crucial enabler for these dynamic scenarios. Changes in job or machine quantities do not require scheduling from scratch; MIStar can perform re-optimization from any valid existing solution [6].

Comparisons with dynamic scheduling algorithms

We appreciate the concern regarding the lack of comparisons with dynamic scheduling algorithms. To address this, we included results of two dynamic methods evaluated on static benchmarks. In future work, we consider to extend MIStar to dynamic scenarios and compare it with more dynamic approaches.

Table: Comparison with Dynamic Algorithms

Reward curve analysis

As we are unable to submit figures, we provide a descriptive analysis of the reward curve observed during training. The training process includes 20,000 episodes, which are grouped into batches of 20 to form 1,000 epochs. We then track the average return per epoch throughout training to assess the model’s learning progress over time.

The curve shows frequent fluctuations in the initial phase (0–200 epochs), ranging approximately between 170 and 210. This high variability reflects the model's active exploration of the policy space. Since new training data is introduced every 20 epochs, the model must continuously adapt to different scheduling instances, which also contributes to these fluctuations. During the mid-phase (200–500 epochs), the curve becomes relatively more stable but still shows several noticeable drops in return, decreasing to around 150–180. This suggests potential local optima or instability in the training process. However, as the model gradually learns a more robust policy, it eventually escapes these suboptimal regions, with the reward converging to a stable range around 230 after approximately 600 epochs.

Ablation on embedding architecture

To accurately extract the key decision-making signals for the policy network—namely, the current scheduling state represented by operation and machine node embeddings—we construct a tailored representation module composed of multiple specialized graph encoders. We adopt GIN[8] to extract structural differences among operation nodes due to its strong discriminative power in graph structure learning. To model the distinct semantic constraints from neighboring nodes (i.e., job predecessors and machine predecessors), we employ a GAT[9] with $n_H$ attention heads, enabling the model to learn the varying importance of different types of neighbors[1][2]. Additionally, we use a HGAT[2] to learn machine embeddings that reflect machine workloads, which serve as critical heuristics for downstream policy decisions—such as guiding the agent to reassign operations from overloaded machines to less loaded ones to reduce makespan.

We conducted ablation studies on the embedding architecture using four size instances from the SD2. Results below further confirm that each of these components is essential for learning informative and discriminative state representations that support effective policy learning.

Robustness

To tackle your robustness concern, we conducted 10 independent runs for each instance in Table 2, and reported the corresponding standard deviation to demonstrate robustness. Due to time constraints, we were unable to include similar statistical measures for the synthetic data in Table 1, but we will incorporate them in the final version to better illustrate performance stability.

Table: Robust Performance across benchmarks（mean ± std）

We are grateful to reviewer for their valuable suggestions. Combined with the discussions and new experimental evidence, these will guide us in revising the final version to boost its clarity and presentation effectiveness.

[1]Zhang C, Cao Z, Song W, et al. Deep reinforcement learning guided improvement heuristic for job shop scheduling[J]. arXiv preprint arXiv:2211.10936, 2022.
[2]Song W, Chen X, Li Q, et al. Flexible job-shop scheduling via graph neural network and deep reinforcement learning[J]. IEEE Transactions on Industrial Informatics, 2022, 19(2): 1600-1610.
[3]Wang R, Wang G, Sun J, et al. Flexible job shop scheduling via dual attention network-based reinforcement learning[J]. IEEE Transactions on Neural Networks and Learning Systems, 2023, 35(3): 3091-3102.
[4]Xie J, Gao L, Peng K, et al. Review on flexible job shop scheduling[J]. IET collaborative intelligent manufacturing, 2019, 1(3): 67-77.
[5]An Y, Chen X, Gao K, et al. Multiobjective flexible job-shop rescheduling with new job insertion and machine preventive maintenance[J]. IEEE Transactions on Cybernetics, 2022, 53(5): 3101-3113.
[6]Renke L, Piplani R, Toro C. A review of dynamic scheduling: context, techniques and prospects[J]. Implementing Industry 4.0: The Model Factory as the Key Enabler for the Future of Manufacturing, 2021: 229-258.
[7]Wu Z, Pan S, Chen F, et al. A comprehensive survey on graph neural networks[J]. IEEE transactions on neural networks and learning systems, 2020, 32(1): 4-24.
[8]Xu K, Hu W, Leskovec J, et al. How powerful are graph neural networks?[J]. arXiv preprint arXiv:1810.00826, 2018.
[9]Veličković P, Cucurull G, Casanova A, et al. Graph attention networks[J]. arXiv preprint arXiv:1710.10903, 2017.
[10]Fuladi S K, Kim C S. Dynamic events in the flexible job-shop scheduling problem: rescheduling with a hybrid metaheuristic algorithm[J]. Algorithms, 2024, 17(4): 142.
[11]Lei, Kun, et al. "Large-scale dynamic scheduling for flexible job-shop with random arrivals of new jobs by hierarchical reinforcement learning." IEEE Transactions on Industrial Informatics 20.1 (2023): 1007-1018.

We appreciate the reviewers' feedback and address each comment below:

Regarding performance, benchmark size, and scalability

We thank the reviewer for these important questions. Our choice of benchmarks, including instances up to 40x10, was intended to align with standard yet challenging problems in the FJSP literature[2][3].We respectfully note that due to the problem's NP-hard nature, these instances are far from trivial.

To the reviewer's point about OR-Tools, its runtime was limited to 30 minutes for experiments in Table 1. The results already show that OR-Tools' ability to find and prove optimality degrades significantly as instance size increases; Due to the complexity of FJSP, even on small-to-medium instances (10×5, 20×5, 15×10, 20×10), OR-Tools fails to find optimal solutions within a reasonable time for most cases. For example, the optimality rate on 20×5, 30×10, and 40×10 instances is 0%. This indicates these are precisely the sizes where powerful learned heuristics become critical.

To more directly address the scalability concern, we conducted new experiments on even larger instances (up to 1,500 operations). For each problem size, 10 instances were randomly generated, and both MIStar and OR-Tools were given a one-hour time limit per instance.

The results, summarized below, highlight two critical advantages of our approach. First and most strikingly, on the highly complex 50x30 instances, OR-Tools failed to produce any feasible solution within the time limit, whereas MIStar consistently found high-quality solutions. This demonstrates MIStar's superior robustness and scalability on truly challenging problems.

Second, for the other large instances where OR-Tools did find a solution, MIStar achieves comparable solution quality in a small fraction of the time. For example, on the 100x10 instances, MIStar reaches a solution with a 66.03% optimality gap in just under 16 minutes, while OR-Tools requires a full hour to achieve a 62.22% gap. This massive speed-up, combined with the strong generalization from a model trained only on smaller 20x10 instances, confirms that MIStar provides an effective and highly scalable solution for large-scale FJSP.

Table: Generalization Performance on Larger-Scale Instances

Adaptation to additional constraints

We thank the reviewer for this important point. We would first like to clarify that our focus on the standard, makespan-oriented FJSP is a deliberate methodological choice. It aligns with the mainstream research paradigm in learning-based scheduling[1,2,3], which prioritizes creating powerful solvers for the core NP-hard problem as a foundational step. Even the static, single-objective FJSP remains a formidable challenge, and developing effective heuristics for it is an active and critical area of research [1,4].

That being said, we fully agree that adaptability to real-world dynamics is crucial. Our core components are highly modular and allow for the incorporation of additional constraints and objectives with minimal, targeted modifications. We are happy to illustrate this adaptability with two representative examples:

Handling machine breakdowns：

Our framework can elegantly manage unexpected events like machine failures. This is possible because our state representation and action space are inherently flexible.

• State: The machine status is a feature in our graph representation, allowing the breakdown to be modeled by updating the corresponding node. The agent, trained to minimize makespan, will naturally avoid assigning operations to unavailable machines or those with infinite processing times, highlighting the effectiveness of our rich, feature-based state representation.
• Action: All operations currently assigned to failed machine must be re-assigned. Broken machines are masked out when generating candidate actions based on the Nopt2 neighborhood structure. As a result, invalid actions $[O_m, O_n, M_k]$ involving broken machines in the third dimension are effectively filtered out. This ensures the agent always operates within valid constraints.
• Reward: To manage the trade-off between efficiency and stability during rescheduling, the reward function can be augmented. A stability term[5] can be incorporated into the reward to minimize the differences in operation start times between the original and revised schedules $\sum_{j=1}^{|J|} \sum_{i=1}^{n_j} \left| st_t - st_{t-1} \right|$, thus encouraging minimal disruption to the former schedule[5,6].

Handling dynamic job arrivals：

MIStar is exceptionally well-suited for dynamic re-optimization, a key advantage of our improvement-based approach over constructive methods.

• State: The flexibility of our graph structure allows new jobs to be seamlessly integrated by expanding the state matrices to include their operations and associated constraints. To maintain focus and efficiency, completed operations can be pruned from the state, ensuring the agent always reasons over the most relevant part of the problem.
• Action: Time-based constraints are easily handled. To respect arrival times, we simply filter out any action that would schedule a job's first operation before it is ready. The problem's inherent precedence constraints, already encoded in our graph, automatically handle the sequencing for all subsequent operations.
• Reward: A penalty term can be added to minimize earliness and tardiness across jobs[7], in order to ensure on-time delivery. Additionally, the robustness term above encourages minimal disruption to the earlier schedule[5].

Finally, the size-agnostic nature of our GNN-based policy network [9] is a crucial enabler for these dynamic scenarios. Changes in job or machine quantities do not require scheduling from scratch; MIStar can perform re-optimization from any valid existing solution [8]. Furthermore, our memory mechanism remains robust, as the incidence matrices from schedules of different sizes can be aligned via simple padding or cropping before similarity comparison. This architectural foresight ensures MIStar is not just a static solver, but a truly adaptable scheduling framework.

Training cost and generalization

MIStar is trained on synthetic data for 20,000 episodes (batch size 20, 100 steps/episode), which takes approximately 4 to 13 days depending on the instance size. A key strength of our approach is that the model does not need to be retrained for different benchmarks. As demonstrated in Tables 1 and 2, our single trained model demonstrates strong scalability and adaptability across benchmarks (Table 2), achieving the highest average performance across instances and generalizing well to larger, unseen instances (30×10 and 40×10 in Table 1).

Regarding the impact of the initial solution

To address the reviewer's question, we have performed a systematic evaluation of MIStar's performance starting from initial solutions of varying quality. This new experiment, detailed in Appendix H (Table 4), uses six different initialization strategies, from DRL-based methods to simple priority rules and random solutions.

The results show that while better initial solutions can lead to higher-quality final schedules within a fixed iteration count, MIStar consistently yields significant improvements regardless of the starting point. Notably, it achieves larger relative improvements when starting from poorer solutions. This demonstrates the robustness of our learned heuristic and its ability to effectively navigate the solution space from diverse initial conditions.

We sincerely appreciate the reviewer’s insightful comments in the rebuttal process. Incorporating these insights, along with the new experimental data and discussions, we will revise the final version of the paper to enhance its clarity and strengthen its overall presentation.

[1]Deep reinforcement learning guided improvement heuristic for job shop scheduling, 2022.
[2]Flexible job-shop scheduling via graph neural network and deep reinforcement learning, 2022.
[3]Flexible job shop scheduling via dual attention network-based reinforcement learning, 2023.
[4]Review on flexible job shop scheduling, 2019.
[5]Dynamic scheduling in flexible job shop systems by considering simultaneously efficiency and stability, 2010.
[6]Robust scheduling for multi-objective flexible job-shop problems with random machine breakdowns, 2013.
[7]Multiobjective flexible job-shop rescheduling with new job insertion and machine preventive maintenance, 2022.
[8]A review of dynamic scheduling: context, techniques and prospects, 2021.
[9] A comprehensive survey on graph neural networks, 2020.

I hope this message finds you well. As the discussion period is nearing its end with less than three days remaining, I want to ensure we have addressed all your concerns satisfactorly. If there are any additional points or feedback you'd like us to consider, please let us know. Your insights are invaluable to us, and we're eager to address any remaining issues to improve our work.

Thank you for your time and effort in reviewing our paper.