Open Materials Generation with Inference-Time Reinforcement Learning
Abstract
Continuous-time generative models for crystalline materials enable inverse materials design by learning to predict stable crystal structures, but incorporating explicit target properties into the generative process remains challenging. Policy-gradient reinforcement learning (RL) provides a principled mechanism for aligning generative models with downstream objectives but typically requires access to the score, which has prevented its application to flow-based models that learn only velocity fields. We introduce Open Materials Generation with Inference-time Reinforcement Learning (OMatG-IRL), a policy-gradient RL framework that operates directly on the learned velocity fields and eliminates the need for the explicit computation of the score. OMatG-IRL leverages stochastic perturbations of the underlying generation dynamics preserving the baseline performance of the pretrained generative model while enabling exploration and policy-gradient estimation at inference time. Using OMatG-IRL, we present the first application of RL to crystal structure prediction (CSP). Our method enables effective reinforcement of an energy-based objective while preserving diversity through composition conditioning, and it achieves performance competitive with score-based RL approaches. Finally, we show that OMatG-IRL can learn time-dependent velocity-annealing schedules, enabling accurate CSP with order-of-magnitude improvements in sampling efficiency and, correspondingly, reduction in generation time. The OMatG-IRL code is included in a new release of the Open Materials Generation (OMatG) framework available at https://github.com/FERMat-ML/OMatG.
Lay Summary
Discovering new crystalline materials is important for technological progress, but the number of possible crystal structures is enormous. AI models can help explore this space by proposing valid and stable candidate structures, but it remains difficult to steer these models toward specific goals, such as a desired property of the crystalline material. We introduce OMatG-IRL, a reinforcement-learning framework that improves crystal generation by giving the model feedback from a target objective and allowing it to learn how to generate better structures. A central feature of OMatG-IRL is that it works even for recent generative models that do not provide a mathematical quantity, called the score, that is usually needed for reinforcement learning. As an example, we apply OMatG-IRL to crystal structure prediction, where the chemical composition is fixed and the model must generate a stable crystal structure for that composition. We use lower energy as the target objective, and the method successfully lowers the predicted energy of generated structures. Importantly, our approach achieves similar performance to the usual reinforcement-learning setup that does use the score. We also show that the method can make crystal generation much faster by learning how to take fewer computational steps during generation, instead of relying on manually tuned rules. These results make reinforcement learning more broadly applicable to generative models and provide a path toward faster, more targeted computational discovery of novel materials.