Open Materials Generation with Inference-Time Reinforcement Learning

TL;DR

This paper introduces OMatG-IRL, a reinforcement learning framework that enhances continuous-time generative models for crystalline materials by enabling property-driven crystal structure prediction without explicit score computation.

Contribution

OMatG-IRL is the first RL application to crystal structure prediction, operating directly on velocity fields, improving sampling efficiency, and maintaining diversity in generative models.

Findings

Effective reinforcement of energy-based objectives in CSP.

Order-of-magnitude improvements in sampling efficiency.

Preservation of diversity through composition conditioning.

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.