Learning Conditional Policy for Crystal Design using Offline Reinforcement Learning

We would like to thank you for your review, and for your constructive comments about the paper. We truly appreciate your points on the clarity of the paper. Please find the explanations for your comments below.

From the experiments, it feels like lacking strong baselines

For our experiments, we primarily focused on unconditional baselines (random, behavioral cloning, and unconditioned CQL) for a better understanding of whether conditioning helps in improving the property of interest. As suggested by Reviewer ffhH, it would also be interesting to compare the performance with conditional baselines with approaches like behavioral cloning (supervised training on materials that have the target property). We are currently working on this aspect.

It would be great if the authors could compare/mention the state-of-the-art methods used in the field of crystal design

Diffusion models[1], GANs[2][3], and RL[4] have been used to address the crystal design problem. Other approaches are cited in Section 2 in the paper. However, as stated in our response to Reviewer b3qf, there are various reasons why a direct comparison of our approach with those models is difficult. An important reason is that most approaches that use DFT for evaluation focus only on a small chemical space or have shape constraints (e.g. cubic). Further, our way of formulating the problem as a sequential decision-making process to predict atom types is entirely different from existing approaches to material design. We do not predict the structure, making it difficult to compare our approach with models that generate the structure. Lastly, none of the previous studies focused on band gap directed inverse design of crystals.

We wish to thank the reviewer for the positive ratings and feedback. If you are satisfied with our response and the updated draft, it would be great if there is any scope for improvement of your rating. If you have more questions, we would be happy to answer them.

References

[1] Xie, Tian, et al. "Crystal diffusion variational autoencoder for periodic material generation." arXiv preprint arXiv:2110.06197 (2021).

[2] Zhao, Yong, et al. "Physics guided deep learning for generative design of crystal materials with symmetry constraints." npj Computational Materials 9.1 (2023): 38.

[3] Zhao, Yong, et al. "High‐throughput discovery of novel cubic crystal materials using deep generative neural networks." Advanced Science 8.20 (2021): 2100566.

[4] Pan, Elton, Christopher Karpovich, and Elsa Olivetti. "Deep reinforcement learning for inverse inorganic materials design." arXiv preprint arXiv:2210.11931 (2022).

We wish to thank you for your review and acknowledgement of our contribution to the underexplored avenue of RL-based inverse design of materials. We have provided explanations for your comments and concerns below.

No novel algorithm is proposed. But for application-oriented research, novel algorithm is not always necessary. The focus here is on implementing and adapting existing methods in a novel context.

We acknowledge and agree with your comment. The primary area of our submission is “applications to physical sciences (physics, chemistry, biology, etc.)”, one of the topics in ICLR 2024 (https://iclr.cc/Conferences/2024/CallForPapers), and while our approach relies on the CQL approach, the construction of our DFT-integrated workflow is unique and novel to the best of our knowledge.

The analyses of the results are limited. It would be beneficial to quantify how many new structures were generated and clarify whether the algorithm simply reproduces known structures from the training set. Given that conservative RL is used, I suspect the output is close to the training set.

The pre-simulation metrics in Section 5.1 of the paper indicate the accuracy and similarity (substitutions with similar atoms) of the trained models on the validation data. While the accuracy, which measures the percentage of accurately predicted atoms globally, was between 36-50% for all models on the validation dataset, the fraction of exactly reconstructed crystals (i.e., all atoms match the ground truth in a crystal) was extremely low (0-7% depending on the model). This could be due to two reasons: 1) the optimization objective of CQL[1] aims not just to fit the dataset (as in BC), but to balance the tradeoff between minimizing the Bellman error and being conservative, 2) the offline policy is trained with batches of transitions i.e., $(s, a, s^{\prime}, r)$, but during test time we perform a rollout starting from an initial state $\boldsymbol{s_0}$. This makes it harder for the learned policy to exactly reconstruct the crystals. Nevertheless, this contributes to the novelty of the generated crystals – more than 90% of the generated crystals of all models had a different composition compared to the corresponding ground truth crystals. Figure 3 in our updated draft shows examples of generated crystals that have the desired property of interest when the corresponding true crystals lie outside the desired range. While the model does reconstruct the crystals in the training set more accurately than the validation set, the points discussed above hold for the training set, too, and this prevents it from reaching the maximum accuracy (100%). About your comment on the extent of the analysis, we have added results for more design targets (2 eV and 3 eV) and added a new metric (out-of-distribution design) in our latest draft of the paper (Figure 2).

Since DFT is actually included in the loop, how about using online RL? Is there any chance to find some novel materials through automated exploration?

Directly using reward signals from DFT calculators in an online reinforcement learning setup is extremely difficult to train because of the large computation time of DFT simulations. However, we believe that the learned offline policy could serve as a good initial policy (thereby incorporating inductive bias) for online RL, and this would be the immediate future work. For online experiments, it is more reasonable to start with a single or a small batch of crystal skeletons in a way that the training process does not take a very long time.

We would like to thank the reviewer for the constructive feedback and positive rating. We hope our responses answered your queries, and are happy to answer more if required. If you are satisfied with the replies and the updated draft, it would be great if there is any scope for improvement in your score for our submission.

References

[1] Kumar, Aviral, et al. "Conservative q-learning for offline reinforcement learning." Advances in Neural Information Processing Systems 33 (2020): 1179-1191.

Thank you for the insightful analysis of the paper. The strengths you mentioned about the reasonableness and usefulness of our methodology and the clarity of the paper are very encouraging. We also appreciate the domain-related comments in your review. Please find our answers to your questions and concerns below.

The policy has a very limited action space, and is basically just learning a more efficient way to substitute atoms.

While this problem is a simplified version of the overall goal of full crystal structure design, we wish to point out that our action space comprises 88 elements in the periodic table (the action space is larger than most of the common RL benchmarks like Atari), and the mapping between the crystal structure and DFT properties is quite complex for any model to learn, which is why we consider this as a ‘learning to search’ problem. Our main hypothesis is the following – among the combinatorial possibilities of atom substitutions given a skeleton of a crystal, there exist some substitutions that result in better stability and have the desired property of interest, and the RL model is expected to search for these solutions. Further, the challenges such as less exploration in offline RL and underestimation of DFT-based band gaps, pose additional difficulties. Also, through this work, we aim to accelerate high-throughput virtual screening (HTVS), which commonly substitutes atoms in existing crystal structures to determine novel crystals – this has resulted in augmentations of popular crystal structure databases such as the Materials Project[1].

I do not see any evidence that the approach could be easily extended/scaled to a more complex action space as described in the future work section, so I am not sure how useful this approach is in practice.

Our future work includes extending the action space to continuous variables like the position of atoms and the lattice structure, in addition to the discrete atom types. This would involve modeling discrete and continuous policies together (or handling multiple action types), which is currently an interesting problem in the RL community. Based on the promising results of previous works that have attempted to tackle this multi-dimensional complex action space, we believe this to be feasible for future work [2][3]. Building on these works for learning policies for atom type and structure generation is a useful extension of our contribution, which addresses the overall problem of designing crystals without prior information about the structure. The inherent challenges of this are yet to be investigated, but we believe that our DFT-integrated offline RL pipeline and evaluation strategies would be very helpful for this challenging design setting.

I would have liked to see more details about how the dataset was generated (e.g., which settings have been used and how were failures handled) if this is one of the main contributions of the paper

Thanks for the suggestion. We have updated our submission with Section A.4 in the Appendix, which explains the details of the DFT computations and handling of failures. As explained in Section 4.4, we construct state trajectories using crystals in the Materials Project dataset; an example of a trajectory is shown in Figure 1a in the paper. Further details are provided in our response to your question on how the dataset was generated.

Total energy says nothing about thermodynamic stability; instead, energy above convex hull is a much better measure.

At present, through our unique reward scheme, we only aim to design crystals with desired band gap that are generally considered stable so they can be used for practical purposes. This is also why the coefficient of the energy term in our reward formulation ($\alpha_1$) is lower than that of the property distance term ($\alpha_2$). However, we understand and acknowledge your concern about total energy. We are constantly looking for reward schemes that better capture the energetic properties of the crystals in a relative sense rather than absolute. While the energy above the convex hull is the best-known measure of a crystal’s stability [4], this metric is conditioned on the elements present in the crystal and is computationally expensive to calculate for a large number of crystals with varying compositions. We are currently redoing our analysis with formation energy, which is not as expensive as energy above convex hull but a more reasonable metric to incorporate in the reward function since it represents a better approximation for crystal stability.

Thank you for your reply.

I agree that substitution is a widely used approach and a more efficient way to do that might be valuable, at least for systems with four or more distinct atom types. My main concern with the future work is not that it is difficult to model, but that it will make the problem significantly harder. In particular, it is much more complex than molecular design since the number of elements is much larger, there are much more degrees of freedom, and there are complex interactions with the lattice.

I disagree that energy above hull is too expensive to compute, but I think formation energy would also be acceptable given that this is an ML conference, so I appreciate your efforts. If you really cannot afford to compute energy above hull for every generated structure, I believe a reasonable compromise would be to use formation energy for the reward but additionally compute energy above hull only on the crystals generated via the validation set. If you could show that your approach still works in this case, I would be willing to revise my score from 3 to 5. However, as it stands I will maintain my current score.

We thank you for the review and feedback, and for acknowledging the importance of the crystal design problem. We appreciate your comment on the usefulness of our method to the material science community. Please find our comments for your questions and concerns below.

This paper is not suitable for the program of ICLR, as it doesn’t advance the state of the art in machine learning.

We wish to clarify that the primary area of our submission is “applications to physical sciences (physics, chemistry, biology, etc.)”, which is mentioned as one of the topics/tracks in ICLR 2024’s call for papers (https://iclr.cc/Conferences/2024/CallForPapers) and therefore represents an important area for ICLR.

CQL is at this point a widely known algorithm and the application to crystal design is not one which pushes ML forward.

We acknowledge that our approach relies on the CQL approach for the offline training. However, to the best of our knowledge, our crystal design pipeline is novel, where we have leveraged large datasets and first-principles DFT calculations, and integrated them into the offline RL training process – while DFT helps in accurate estimation of properties like energy and band gap, the cost of computation is very high. It takes a few seconds to several minutes (~2 minutes on average) to compute the reward for each input crystal. Through our approach, we demonstrate promising results for band-gap-directed crystal design, highlight the challenges and scope for future work, and propose to release a benchmark dataset for the benefit of the material discovery community. We believe that this is a solid contribution to the “applications to physical sciences” track in ICLR.

It is also not clear to me that the MDP formalism is the correct one for this problem setting. It seems to me that using a conditional generative model over the entire data would be a more natural choice (think diffusion model or autoregressive transformer). The authors don’t evaluate against those baselines which I think might be better suited.

Diffusion models have previously been studied for material design[1]. However, this involves predicting the identities/properties of all atoms at once with no sequential process involved. The case is the same for other generative models like GAN. We believe that by framing this problem as a sequential decision-making process, the important aspects of exploration, credit assignment, and ‘learning to search solutions’ are more intuitive, especially when integrated with DFT – this could help in improving the explainability of the model. MDP formalism has proven very useful in small molecule discovery (drug design)[5][6], and is also of interest to the material discovery community[7]. We primarily focused on unconditional baselines (random, behavioral cloning, and unconditioned CQL) in our experiments for an understanding of whether conditioning helps in improving the property of interest. We are currently working on evaluating diffusion and decision transformer-based conditional models for band gap directed crystal design.

Have people tried generative modeling approaches to this problem? How well did they work? Can we quantitatively compare?

CDVAE[1] is the first approach to address the crystal structure design problem using a diffusion model. However, there is no integration of first-principles DFT properties, and only approximate metrics and properties were used for performance evaluation. Other works focused on generative models like GANs, with the most recent one being [2], where generated crystals were evaluated using DFT. Another work is [3], which used a distributional graphormer model to generate carbon-based crystal structures. While the performances of these methods are promising, the reasons why we cannot perform a direct comparison of these models with our approach are mentioned below.

• Many approaches that use DFT focus only on small chemical systems, with lesser number of elements in the vocabulary, or have constraints on the shape of crystals. For instance, Distributional graphormer[3] generates only carbon-based crystal structures. CubicGAN[4] focuses only on cubic crystals. However, a conditional variant of the GAN model proposed by [2] could be more suitable for comparison with our model in terms of property-directed generation. We are working on the same.
• Some of the metrics used by structure generation models (e.g., structural validity as in [1]) cannot be directly used for comparison with our model, since we do not generate the structure.
• None of the previous studies focused on inverse design of crystals (without structural or composition constraints) with band gap as the property of interest.