Open Materials Generation with Inference-Time Reinforcement Learning

TL;DR

Open Materials Generation with Inference-time Reinforcement Learning (OMatG-IRL) addresses the challenge of aligning crystal structure generation with explicit target properties when using flow-based models that may not expose a score. The method performs policy-gradient reinforcement learning at inference time directly on learned velocity fields within the stochastic interpolants framework, with variants for score-based and velocity-based updates and a learned velocity-annealing schedule. Key contributions include the first RL application to CSP, showing energy-based rewards can be effectively optimized without explicit diversity rewards, and achieving substantial gains in sampling efficiency by reducing integration steps by up to an order of magnitude. The results demonstrate that velocity-based RL can match score-based RL, broadening the applicability of RL to a broader class of generative models for materials design. Together, these advances enable more efficient, diverse, and target-driven crystal structure generation.

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.

Figures (8)

• Figure 1: Inference-time RL for CSP in velocity-based OMatG-IRL. (a) The deterministic base ODE with pretrained velocity field $b^{\theta_\text{ref}}(t,x_t)$ is augmented with a small noise schedule $\sigma_\text{ref}(t)$, yielding a surrogate SDE that leaves evaluation metrics (e.g., deviation from a reference structure) of the final samples $x_{t=1}$ virtually unchanged. (b) The frozen surrogate defines a reference policy for KL regularization, while stochastic exploration is performed using a reinforced velocity field $b^\theta(t,x_t)$ and a (potentially different) noise schedule $\sigma(t)$. (c) GRPO compares terminal rewards $r^i=r(x^i_{t=1})$ obtained from multiple stochastic-policy rollouts under identical conditioning. (d) These rewards are transformed into GRPO group advantages and used, together with KL regularization, in a PPO-style clipped objective to update the policy.
• Figure 2: Test-set evaluation metrics for score-based SDE and perturbed velocity-based ODE integration of the atomic positions under different noise schedules. Small, medium, and large noise scales are denoted by $a_s$, $a_m$, and $a_l$, respectively.
• Figure 3: Evolution of validation metrics for score-based and velocity-based OMatG-IRL as a function of RL training iteration, shown for three random seeds of the same setup ($N_t=50$). The colored dashed lines indicate the test-set performance of the OMatG-IRL checkpoint selected by the validation optimum. For reference, we also show the test-set performance of the original velocity-annealed OMatG model evaluated with $N_t=740$ integration steps. The velocity-based OMatG-IRL setup with small reference noise scale $a_s$ (green) uses fewer PPO epochs per update, resulting in slower but more stable reinforcement.
• Figure 4: Test-set evaluation metrics for velocity-annealing OMatG and OMatG-IRL as a function of the number of integration steps $N_t$, highlighting the improved robustness of OMatG-IRL to aggressive time discretization.
• Figure 5: Velocity-annealing schedules for the atomic positions and lattice vectors, either learned with OMatG-IRL or obtained from a hyperparameter sweep over a handcrafted schedule at a given number of integration steps. The learned schedules adapt qualitatively differently from handcrafted ones.