CrystalGRPO: Target-Aligned and Coverage-Preserving Reinforcement Learning for Flow-Based Crystal Structure Prediction

2026-08-06

AI Summary

This work addresses limitations in existing flow-based crystal generation models, which fail to directly optimize target structure recovery during pretraining and rely solely on energy-based rewards and random coordinate perturbations in reinforcement learning fine-tuning, often compromising polymorph identification accuracy and structural diversity. To overcome these issues, the authors propose CrystalGRPO, a novel framework that extends the ODE-to-SDE strategy to the joint coordinate–lattice space for the first time. Integrating MACE-based energy prediction with StructureMatcher-derived recovery scores, CrystalGRPO introduces two variants: CrystalGRPO-Q, which optimizes single-sample recovery, and CrystalGRPO-C, which incorporates trajectory regularization and coverage-aware advantage estimation. Evaluated on MP-20 and MPTS-52 benchmarks, CrystalGRPO-Q significantly improves Top-1 recovery accuracy, while CrystalGRPO-C enhances Top-20 success rates and reduces RMSE, effectively balancing precision and diversity.

Abstract

Flow-based generative models can efficiently produce candidate structures for crystal structure prediction (CSP), but their pretrained objectives do not directly optimize downstream target recovery. Reinforcement-learning post-training offers a flexible solution, yet existing approaches rely primarily on energy rewards and coordinate-only stochastic policies. Predicted energy does not identify the reference polymorph, while reward-driven concentration can reduce the candidate coverage required for Top-N recovery. We introduce CrystalGRPO, a CSP-aligned post-training framework that extends existing ODE-to-SDE policy constructions to the joint coordinate--lattice state. CrystalGRPO combines MACE-predicted energy with a StructureMatcher-based recovery score and provides two operating modes: CrystalGRPO-Q, which prioritizes single-draw recovery, and CrystalGRPO-C, which combines full-trajectory reference regularization with a coverage-aware group advantage to preserve finite-budget target recovery. Across MP-20 and MPTS-52 with PXRDGen and OMatG backbones, both variants reduce one- and twenty-sample RMSE relative to coordinate-only reinforcement in all four backbone--dataset settings. CrystalGRPO-Q consistently improves Top-1, whereas CrystalGRPO-C achieves a higher Top-20 across all settings.

Problem

• crystal structure prediction
• reinforcement learning
• flow-based generative models
• target recovery
• candidate coverage

Innovation

• CrystalGRPO
• flow-based generative model
• reinforcement learning
• crystal structure prediction
• coverage-preserving