MACS: Multi-Agent Reinforcement Learning for Optimization of Crystal Structures

Dear Reviewer,

Thank you for taking the time to read our paper and provide valuable feedback.

Weakness: Although the paper casts atoms as agents in a MARL setting, the agent design lacks meaningful interaction or interdependence typically expected in multi-agent systems. There is no explicit coordination mechanism beyond standard parameter sharing. This setup resembles multiple parallel optimizers more than a principled MARL approach.

Weakness: The “agent” in MACS functions more like a local function approximator on force minimization than a true autonomous or strategic decision-maker in a multi-agent setting. The use of the MARL framework here feels superficial and may be misleading in terms of the novelty claimed.

Answer to Weakness 1. We would like to refer the reviewer to our Summary Rebuttal (at the beginning of the Rebuttal to ytGL) that addresses these questions.

Weakness: The paper compares against standard MARL baselines, but does not include strong compositional or hierarchical RL baselines

Answer to Weakness 2. To the best of our knowledge, we are the first to apply MARL to address periodic crystal structure optimization (as mentioned in Lines 58-59 in our paper). All baselines are not MARL approaches, and we are not aware of any existing hierarchical RL methods used for this task. We believe hierarchical RL methods could be more suitable for solving global structure optimization, which is an interesting direction for future work.

Weakness: No theoretical insights or guarantees are provided on the benefits of compositionality, generalization bounds, or conditions under which MACS would outperform flat policies.

Answer to Weakness 3. We appreciate the reviewer’s point and agree that providing formal theoretical guarantees for performance or generalization in this setting remains an open challenge. Our primary contribution focuses on designing and validating an MARL-based framework that addresses the complexity of crystal structure optimization, which involves highly non-convex, high-dimensional energy landscapes with strong physical constraints. Deriving theoretical guarantees in such environments is non-trivial because reward functions and state spaces depend on material-specific physics rather than the synthetic or formal assumptions typically used in theoretical analysis.

That said, our results provide strong empirical evidence that the MARL framework preserves transferability within a phase field and maintains scalability for structures twice as large as the structures the policy was trained on. We agree that establishing theoretical conditions for guarantees on optimization time and quality - whether for existing methods or new ones - would be a significant achievement, which we leave for future research.

Dear Reviewer,

Thank you for taking the time to read our paper and provide valuable feedback. In the following, we address your concerns and questions in detail.

Concern 1: To me, the connection to learning to optimize is somewhat superficial: the RL algorithm is standard, the environment (CHGNet + atomic displacements) is fixed, and there is no higher-level learning over tasks.

Answer 1. Indeed, MACS is tailored specifically for geometry optimization and does not support optimization tasks from other areas of science. In the manuscript, we note that "an entire field within materials design, known as crystal structure prediction (CSP), focuses on the computational prediction of stable crystal structures" thus geometry optimization plays a fundamental role in computational materials design. Due to the highly diverse chemistry across different compositions and phase fields, we treat optimization within each composition or phase field as an individual task over which MACS learns to optimize.

Concern 2: The agents operate independently (using independent SAC). While effective, this ignores coordination issues known to affect MARL systems. I believe a short discussion would be helpful to better understand how strongly these issues might affect the proposed Markov game.

Answer 2. Summary Rebuttal (at the beginning of the Rebuttal to ytGL) partly addresses this question. We would like to add that previous works [4, 5] have also shown that independent learning methods can perform just as well as or even better than state-of-the-art joint learning MARL approaches on multi-agent benchmark environments like the StarCraft Multi-Agent Challenge (SMAC), which consists of various challenging coordination tasks.

Further, the introduced action scaling coefficient $c_i$, defined by formula (3), restricts the order of magnitude of an agent's action based on its gradient norm. This makes transitions between states more predictable for agents, as each agent can rely on the gradients of neighboring agents in its local observation to estimate how far they can move from their current positions. Together with overlapping local observations and interdependent rewards, this enables efficient collective behavior without explicit coordination.

[4] De Witt, C.S., Gupta, T., Makoviichuk, D., Makoviychuk, V., Torr, P.H., Sun, M. and Whiteson, S.. Is independent learning all you need in the starcraft multi-agent challenge? arXiv preprint arXiv:2011.09533, 2020.

[5] Papoudakis, G., Christianos, F., Schäfer, L. and Albrecht, S.V.. Benchmarking multi-agent deep reinforcement learning algorithms in cooperative tasks. In Neural Information Processing Systems Track on Datasets and Benchmarks, 2021.

Concern 3: My main concern is that, although the results look strong, the baselines are limited to classical optimizers (e.g., BFGS, FIRE). There has been a number of works on geometric optimization, e.g., [1-3]. I believe the confidence in MACS' performance would be substantially improved if stronger baselines were to be included.

Question 2: The confidence in MACS' performance would be substantially improved if stronger baselines were to be included.

Answer 3. We refer the reviewer to the Summary Rebuttal to answer the question. Also, in the papers [1-3], the structure optimization method provided is either designed for a different task or tailored to more restrictive conditions than those addressed in our paper. Specifically, in [1], the optimization method is applied under symmetry constraints to crystal structures close to local energy minima; in [2], the authors perform two-step structure optimization using two different energy calculators and classical optimizers; in [3], the authors consider the specific problem of structure optimization under noise in gradient conditions and introduce a two-staged optimization via the steepest-descent method with a fixed step size on each phase. The steepest descent is known to have slower convergence than other first-order methods, such as CG, one of the baselines in our study.

Concern 4: The energy estimations are based on CHGNet, which is itself a learned model. While fast, it introduces approximation errors. Through exploration, the RL agents might produce crystals that are o.o.d. of CHGNet's training data. It would be valuable to confirm the quality of the final crystal structures using DFT.

Answer 4. The final energy distributions shown in Fig.12 of Appendix B.4 confirm the comparable final energy distribution among results of the optimization by MACS and the baselines. Performing density functional theory calculations on all of the final structures is impractical within the available time scale. However, we do agree with the reviewer, that assessing the quality of the resulting final crystal structures is an important task. For this we have performed an additional analysis of the final structures and suggest to add the following section to the Appendix:

"The analysis of structures optimized by MACS.

We calculated the atom-atom distances up to 6.00 Angstroms within all 7,200 structures which were optimized using MACS. For each composition, we tabulated the distances, rounded to 2 decimal places as histograms, all of the histograms for individual compositions were then summed to give a total atom-atom distance histogram for each composition. For each composition, we have then tabulated the shortest distance observed, and the distance of the first major peak in the histogram, both provided in Angstroms. We provide these results below with the total number of counts for the distances observed in brackets.

---

The shortest interatomic distances observed are plausible distances, the major observed peaks for each composition correspond to expected values when considering each chemistry. For example, the first peak at 1.87 Angstroms in SrTiO$_3$, is close to that of the Ti-O distances in the experimentally observed structure of 1.95 Angstroms. These observations in addition to the distribution of energies shown in Fig.12 of Appendix B.4 lead us to the conclusion that the final structures produced by MACS do not contain unphysically short interatomic distances that are typical for structures produced due to the failure of the energy calculator."

Dear Reviewer,

Thank you for all of your kind suggestions and questions. In the following, we answer all of your questions in detail.

Weakness (1): The novelty of the paper is limited. The primary contribution of the paper is formulating the optimization of crystal structure as a POMDP by defining an appropriate state and action space and reward function. Otherwise, the algorithm is a standard application of SAC. Additionally, while the paper refers to multi-agent RL, the actual algorithm treats each agent as a independent agent optimizing indepedent rewards. While an important and interesting application of RL, the core machine learning contribution is limited.

Answer to weakness (1): We would like to refer the reviewer to the Summary Rebuttal (at the beginning of the Rebuttal to ytGL) for the broad discussion. We acknowledge that exploring richer multi-agent interactions (e.g., via communication or centralized training) is a possible direction for future work.

Weakness 2: It is not clear to me if the baseline comparisons are entirely fair. For example - are the reported metrics and inclusive of the training steps required to train MACS? Finally MACS would produce stochastic policies but I do not see a discussion of the variability in policy performance across seeds.

Question (1): I am not 100% clear if the training computation is included in the metrics for comparing to baselines. My understanding is that it is not. Is this a fair comparison?..

Answer to Question (1): Indeed, the training time is not included in the comparison with the baselines. We suggest explicitly mentioning it in the main text and referring to a new section in the Appendix for the detailed discussion. The proposed text for that section is as follows:

"Appendix Section: Training Time Contribution to the Experiments.

The training time of MACS is 343,800 seconds and was not included in the evaluation of MACS. We evaluate the MACS policy based on optimization time only because including the combined training and testing time would make the evaluation highly sensitive to the size and number of test sets. In fact, the experiments could be extended by increasing the sizes/numbers of the test sets until the relative contribution of training time becomes negligible, leading to the reported results.

Moreover, this setup reflects practical usage, as screening hundreds of thousands of structures is a typical use case in computational chemistry. Furthermore, the scalability and zero-shot transferability of MACS allow it to be trained on one set of compositions and then used to efficiently optimize larger structures from other compositions, without additional training, i.e., with zero training time in those cases.

There is also no fair way to estimate the combined training and testing time, since these phases were conducted under different conditions. Although the same hardware was used the difference in the number of active cores (see Appendix B.1) raises the question of whether to report runtime using wall-clock time or CPU time, and which approach is fairer.

For reference, comparing the combined training and optimization time of MACS with the optimization time of the fastest baseline (BFGSLS), MACS is 3% faster in terms of wall-clock time but 126 times slower in terms of CPU time. Given the fact that training time can be omitted when the method is used in practice, we compare the methods only by the optimization time, as shown in Table 1."

Question (2): Failure rate is not defined and its positioning in Table 1 is a bit confusing. Is this the overall failure rate across all composition? Can this be reported per composition instead?

Answer to Question (2). The failure rate notation $\mathbf{P}_{F}$ is introduced on line 240. The failure rate reported in Table 1 represents the average across all test sets. Since the average failure rate is low and comparable to the best baseline, we provide failure rates per test set in Table 7 of Appendix B.4 and refer to it in the caption of Table 1. Table 7 shows that the failure rate remains consistently low across all test sets. We will add a note to Table 1 to clarify that the numbers provided there are averaged across the test sets to avoid the confusion.

(3) SAC produces a stochastic policy, however I do not see any discussion of evaulation over multiple seeds - was this done?

Answer to Question (3): We initially did not evaluate MACS on the same structures across multiple seeds, as we relied on a large number of randomly generated structures (300 per test set, across 24 test sets) to ensure reliable policy evaluation. Following the reviewer’s suggestion, we have conducted an additional experiment to assess variability due to the stochastic nature of SAC used by MACS. We propose adding the following section to the Appendix and referencing it in the main text:

"Appendix Section: Variability of the MACS Policy. We evaluated the policy on a test set containing SrTiO$_3$ structures with 40 atoms by performing nine additional optimization runs, resulting in ten runs in total for this set. We measured variability in the average number of steps for successful optimizations, and the failure rate.

The failure rate was 0.33% in 2 runs, 0.66% in 1 run, and 0% in all others, confirming that failures remain rare.

For the number of optimization steps/energy calculations, we estimated (mean, std) among 10 seeds as (147.48, 3.18) which is well below the mean number of steps (190) in the BFGSLS, the best baseline in this test set."

Question (4): In line 295 it states that episodic reward is normalized to compare across reward functions - normalized in what way?

Answer to Question (4). All episodic rewards collected during training are normalized using min-max (unity-based) normalization to fit within the range $[0,1]$. This normalization avoids comparisons based on absolute reward values and instead allows assessment of convergence speed and stability.

Limitations: The authors address several limitiations, however I would add that the claim in line 319 seems too strong for the results provided in the paper since the test cases only consist of structures that contain the same element compositions as the training cases.

Answer to Limitations. For any phase field with at least two elements, the number of possible compositions is infinite. In this context, the ability of MACS to apply a policy trained on one composition to optimize all structures within the same phase field makes it a truly universal optimizer for that field. To enable transfer learning across different phase fields, additional training on compositions from diverse phase fields will be required. We suggest to replace the sentence in line 319 by the following sentence: "In conclusion, MACS has the potential to evolve into a universal geometric optimizer for periodic crystal structures."

Summary Rebuttal

To avoid repetition, we have formulated a shared response for the reviewers addressing the questions and concerns raised by more than one reviewer.

We believe that simplicity is a strength and have intentionally used the simplest MARL approach with independent learning. Instead of coordination via explicit communication, our method enables coordinated behavior through interactions among agents via interdependent rewards and overlapping local observations. This allows agents to adapt to each other’s behavior over time, resulting in implicit coordination. Another main reason is that independent methods tend to scale better to tasks involving a large number of agents and/or actions compared to centralized methods. This is crucial for our method’s scalability and applicability in computational chemistry, where the optimization of large structures is a major bottleneck in many computational methods.

Our experiments demonstrate that, with a carefully designed setup of actions, observations, and rewards, independent learning can offer a significant improvement over state-of-the-art methods for crystal structure geometry optimization. We recognize the importance of convincing you that our baselines are strong enough for this claim to be valid, and we address this below.

The baselines we use, although classical, still form the majority of methods available in widely used packages for computational chemistry, such as ASE, VASP, GULP, and CHGNet. These are employed in thousands of computational research projects. In short, our baselines represent the standard choices used daily by practitioners working in computational chemistry. The lack of novel methods that consistently outperform these baselines in general settings motivated our research. We appreciate the references provided by Reviewer 6w1f; however, as we explain below in our rebuttal to 6w1f, none of these are appropriate for our task.

Finally, we would like to emphasize the contribution of our paper in formulating a challenging, physically grounded optimization environment for a ubiquitous task in condensed matter science, and in proposing an efficient, scalable MARL method that opens new directions for further research.

Rebuttal to ytGL

Dear Reviewer,

Thank you for taking the time to read our paper and provide valuable feedback. We will fix the typos and improve Table 1 and Figure 3 as suggested. The answers to your questions are below.

Question 1: Did you tune BFGS-LS and CG for the experiments?

Answer 1. All baselines were used with the default hyperparameters defined in the ASE package, as they represent a typical use case.

Question 2: Any preliminary results on treating lattice vectors as extra agents, as suggested in Limitations?

Answer 2. This is an interesting future work, as mentioned in our paper, and preliminary results are not yet available.

Limitations: The paper notes that atomic descriptors are minimal and unit-cell vectors are fixed/ it could also discuss how reliance on CHGNet accuracy affects performance and how the policy behaves if force evaluations are noisy.

Answer 3. We have conducted the additional analysis of the quality of the final structures produced by MACS, which is provided in the reply to reviewer 6w1f (Answer 4) and will be added to the Appendix. This analysis of the resulting structures reveals that they are physically acceptable. The noisy force evaluations would be an interesting direction for future research, and we will mention it in the revised paper.