Reinforcement Learning on the Discrete Composition Channel of a Crystal Generator: Validated Gains and Reward Hacking

Pawan Prakash; Philipp Höllmer; Addis Fuhr; Peter Hirschfeld; P. Ganesh; Stefano Martiniani; Richard Hennig

cs.LGOct 2, 2026

Reinforcement Learning on the Discrete Composition Channel of a Crystal Generator: Validated Gains and Reward Hacking

Authors: Pawan Prakash, Philipp Höllmer, Addis Fuhr, Peter Hirschfeld, P. Ganesh, Stefano Martiniani, Richard Hennig

Organizations: University of Florida · New York University · Oak Ridge National Laboratory

Abstract

Inverse materials design is a long-standing goal of computational materials discovery. Generative models for crystalline materials are typically trained to match the distribution of a structure database, while nothing in their training objective points them at specific design goals such as targeted properties. We use group-relative policy optimization (GRPO) to align a generative model based on stochastic interpolants and discrete flow matching with general black-box reward functions through reinforcement learning. Atom types are generated by a discrete flow and the policy gradient of our generalization of GRPO directly acts on the likelihoods of the atom-type transitions, which differentiates our work from previous reinforcement-learning approaches for diffusion and flow-based generative models of crystalline materials. We introduce a reward function that raises the yield of metastable, unique and novel structures (mSUN) from 13.4% for the pretrained model to 45.5% for the reinforced model, as evaluated by a community benchmark. Our reward also improves the performance of a reinforcement learning framework for crystalline materials based on latent denoising diffusion models. At the same time, we find that directly reinforcing atom-type transition likelihoods enables reward exploitation that has to be prevented with explicit guards. The same analysis also exposes a gap in the community metric. Single-element structures in distinct packings are counted as metastable, unique and novel materials and inflate mSUN without yielding any new compounds. A stability claim is only as good as its reference hull. We report every result split by the number of reference phases behind it and argue that benchmarks should do the same.

Figures & tables

Figure 1: Reinforcement learning on a crystal generator can satisfy the benchmark without real discoveries, but guards in the reward close the exploit. (a) The OMatG generator turns noise with masked element identities into crystals and is scored by mSUN under LeMat-GenBench (Section 4.1 ). Without the single-element guard, RL inflates the pretrained model’s 13.4% to 37.6%, but over half of that mSUN set is single-element repackings of known elements, such as face-centered-cubic Mg counted as a novel material. With the full reward (right, Section 3.3 ), which routes sparse-hull structures to abstention and single-element structures to the penalty reward, mSUN reaches 45.5% with 0.4% single-element structures and DFT-confirmed compounds on well-referenced hulls. (b) The composition channel. Under masked discrete flow matching, each atom’s element identity is a discrete commit event with an exact log-probability, while positions and lattice evolve as a continuous SDE (Section 3.2 ).

method	valid	unique	novel	metastable	SUN	mSUN	single-element	sparse hull
	% $\uparrow$	% $\uparrow$	% $\uparrow$	% $\uparrow$	count $\uparrow$	% (count) $\uparrow$	in mSUN $\downarrow$	% of mSUN $\downarrow$
prior	94.1	99.5	82.1	27.2	8	13.4 (336)	7	21.7
best-of- $N$	97.9	77.9	11.5	65.6	72	7.7 (192)	0	0.0
frozen control	93.5	99.5	87.7	18.7	9	9.8 (245)	1	12.2
Chemeleon2	94.9	90.6	68.1	67.5	36	36.6 (915)	9	7.9
OMatGRPO (ours)	69.8	98.7	88.3	75.8	7	45.5 (1138)	4	23.5

Table 1: Yield of mSUN under LeMat-GenBench. Valid is the share of the nominal 2,500 generations that passes the benchmark’s structural checks. Unique, novel and metastable are rates over valid structures, the benchmark’s own convention, and mSUN is reported over the nominal 2,500 so that methods are comparable (counts in parentheses). SUN and single-element mSUN are counts. Prior is the pretrained OMatG model. The zero single-element and sparse-hull entries of best-of-

N

are enforced by its selection score rather than learned, and the controls are defined in Section 4.2 .

Figure 2: Reinforcement learning gains, their controls, and DFT validation. Every method submits 2,500 generated structures, relaxed and scored by LeMat-GenBench (Section 4.1 ). Counts are printed above the bars. (a) Solid segments sit on trusted hulls with at least 12 near-hull references, lighter segments of the same color on sparse hulls (Section 3.3 ), and error bars are 95% confidence intervals. The controls are defined in Section 4.2 . (b) Compounds only, which removes the single-element structures of Section 4.3 . (c) MLIP against DFT stability claims for the spot-checked structures. On trusted hulls (

n{=}67

) the mean absolute deviation is 0.016 eV/atom and 63 of 67 claims agree. On sparse hulls (

n{=}21

) the energies track but stability cannot be adjudicated. Appendix E lists every structure, the four disagreements and the two excluded cases.

Figure 3: Reward hacking via single-element structures. Throughout, discovery is the reward of Eq. 5 with the single-element guard removed. a , Composition of each method’s mSUN set, grouped by the number of distinct elements per structure, with the pretrained prior and Chemeleon2 for reference. b , Fraction of generations scored on sparse hulls and fraction that are single-element during training, smoothed over 25 rollout iterations, for the discovery run and OMatGRPO.

	trained in pipeline
reward	ours (native-channel SDE)	Chemeleon2 (latent VAE)
ours	45.5 (1138)	43.7 (1093)
Chemeleon2	18.6 (465)	36.6 (915)

Table 2: Swapping rewards between the two pipelines. Each cell is a policy trained with the row’s reward inside the column’s pipeline and evaluated as in Table 1 , given as mSUN in % of the nominal 2,500 with counts in parentheses. The diagonal holds OMatGRPO and the released Chemeleon2 checkpoint. Each cell keeps its pipeline’s relaxation convention, ours relaxes before the reward and Chemeleon2’s does not, and Appendix D reports both variants.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

	internal evaluation	LeMat-GenBench
sample per method	$1{,}000$ generations	$2{,}500$ generations
rate basis	count per $1{,}000$ generations	fraction of the nominal $2{,}500$ (mSUN) or of the valid structures (other rates)
relaxation before scoring	FIRE under UMA with cell, at most 500 steps	our relaxed geometry, then the benchmark’s own 50-step relaxation
stability model	UMA	Orb, MACE and UMA, each against its own hull, two of three must agree
reference hull	UMA-consistent LeMat-Bulk hull	the benchmark’s per-potential hulls
metastability threshold	$E_{\mathrm{hull}}\leq 0.1$ eV/atom	$E_{\mathrm{hull}}\leq 0.1$ eV/atom

Appendix

Table 3: The internal evaluation and LeMat-GenBench side by side.

run	identifier	mSUN	%	single-element	sparse hull	compounds
prior	R0	336	13.4	7 (2.1%)	21.7%	329
best-of- $N$	bestofn48k_top2500	192	7.7	0 (0.0%)	0.0%	192
frozen-composition control	frozen_control	245	9.8	1 (0.4%)	12.2%	244
Chemeleon2, released model	chemeleon2_mp20rl	915	36.6	9 (1.0%)	7.9%	906
discovery, pre-creativity	canonical	1063	42.5	528 (49.7%)	65.8%	535
discovery (Section 4.3 )	canonical_creatrelax	939	37.6	502 (53.5%)	63.4%	437

Appendix

Table 4: LeMat-GenBench results of every run named in the paper, on the nominal

2{,}500

basis. Single-element and sparse-hull entries are shares of the run’s mSUN set, and compounds is the mSUN count without the single-element structures. The identifier is the run name in the released code and data, without the suffix _r750 . The five runs of Table 1 are repeated for reference.

#	reward optimized	what the policy did	why it paid	closure
1	relaxation displacement (RMSD)	near-stationary but energetically implausible structures	displacement carries no energy information	energy-based reward
2	formation energy	collapse onto deeply bound, far-above-hull fluorides	strong bonding is a poor proxy for phase stability when more stable competitors exist	switch objective to $E_{\mathrm{hull}}$
3	unfloored $E_{\mathrm{hull}}$	below-hull depth chasing in sparsely referenced regions	with few references $E_{\mathrm{hull}}$ is unbounded below, and depth pays linearly	floor at zero. $-\mathrm{clip}(E_{\mathrm{hull}},0,1)$ removes the payout by construction
4	floored $E_{\mathrm{hull}}$ $+$ residual geometry term	sparse above-hull O/F binaries gaming the geometry term	thinly referenced hulls score leniently	route sparse-reference scores away from credit
5	full recipe, abstention routing	distinct single-element packings flooding the mSUN set	single-element subsystems are sparsely charted, and the packings pass the metastability, uniqueness and novelty checks of the three community implementations we checked at source	single-element guard. A single-element composition receives the penalty reward, and no relocation followed

Appendix

Table 5: The reward-hacking taxonomy. Five specification-gaming modes met under displacement, energy and

E_{\mathrm{hull}}

rewards, each closed by one targeted change.

run	single-element in mSUN	note
pretrained prior (MP-20)	6
discovery recipe, pre-creativity, $K{=}16$	253	abstention routing, no single-element guard
discovery recipe, pre-creativity, $K{=}32$	307	matched rollout budget
$+$ single-element guard, pre-creativity	12	2 to 15 across three seeds
discovery run (Eq. 5 without the guard)	242
OMatGRPO (Eq. 5 )	2

Appendix

Table 6: Single-element structures in the mSUN set. Upper block, per

1{,}000

generations under the internal evaluation. Lower block, raw counts in released sample sets of generators that were sampled rather than trained against the metric.

Figure 4: Where the single-element structures come from. Single-element structures in the mSUN set per

1{,}000

generations under the internal evaluation. On the MP-20 prior the guard reduces the count from 242 to 2. The same training induces these structures from zero on a second model pretrained on Alexandria, and on the MP-20 model a larger group size at the same rollout budget produces more of them. The Alexandria and group-size comparisons were run before the creativity term was added to the reward.

run	$n_{\text{val}}$	valid	uniq	novel	meta	SUN	count	/val	/2500
		rate on valid					mSUN
prior (no RL)	2352	0.948	0.995	0.821	0.272	8	336	0.143	0.134
ours $\times$ ours	1745	0.703	0.987	0.883	0.758	7	1138	0.652	0.455
theirs $\times$ theirs	2372	0.949	0.906	0.681	0.675	36	915	0.386	0.366
B (theirs $\times$ ours)	2305	0.928	0.990	0.812	0.321	3	401	0.174	0.160
B ′ (B $+$ relax)	2322	0.933	0.992	0.817	0.341	6	465	0.200	0.186

Appendix

Table 7: Full per-run LeMat-GenBench results. Rate columns are on the valid set (

n_{\text{val}}

), and valid itself is on the scored set. mSUN is given as count, rate on valid, and rate on the

2{,}500

nominal draw.

run	single-element in mSUN	sparse in mSUN	dominant arity	$\overline{E_{\mathrm{hull}}}$	$\overline{E_{\mathrm{form}}}$	relax RMSE
prior (no RL)	2.1%	21.7%	3 (50%)	0.192	$-0.733$	0.441
ours $\times$ ours	0.4%	23.5%	3 (46%)	0.070	$-0.055$	0.278
theirs $\times$ theirs	1.0%	7.9%	3 (48%)	0.078	$-0.804$	0.168
B	7.5%	26.7%	3 (44%)	0.178	$-0.735$	0.456
B ′	4.1%	22.2%	3 (45%)	0.161	$-0.689$	0.448
A	0.0%	15.3%	3 (49%)	0.143	$-0.779$	0.379

Appendix

Table 8: Compositional and energetic companions per run. Single-element and sparse-hull shares are over each run’s mSUN set, and dominant arity is the most common element count and its share. Means over the valid set, energies in eV/atom, relaxation RMSE in Å.

term	untrained policy	trained policy	$\Delta$
creativity	1.000	0.994	$\approx 0$
composition diversity	0.473	0.648	$+0.175$
energy	0.043	0.125	$+0.082$
structure diversity	0.035	0.044	$+0.009$
total	1.551	1.811	$+0.260$

Appendix

Table 9: Reward composition of the B run under the untrained and trained policy (post-normalization means).

tier	run	$n$	MAE (eV/atom)	agreement
trusted	prior	4	0.005	3 of 4
trusted	discovery, pre-creativity	28	0.009	28 of 28
trusted	guarded, pre-creativity	13	0.019	12 of 13
trusted	OMatGRPO	22	0.026	20 of 22
trusted	all	67	0.016	63 of 67
sparse	discovery, pre-creativity	10	0.033	8 of 10

Appendix

Table 10: UMA against DFT hull distances by tier and by run. Agreement counts metastability claims on the same side of 0.1 eV/atom under both methods.

run	formula	tier	refs	UMA	DFT	agree
D -	Ac 3 Hg 5	T	9	0.054	0.050	y
D -	Ac 9 In 3	T	6	0.080	0.084	y
D -	AgLi 8	T	11	0.007	0.013	y
D -	AgLi 9	T	11	0.040	0.043	y
D -	AgMg 9	T	7	0.068	0.085	y
D -	AgSr 9	T	12	0.088	0.024	y

Appendix

Table 11: Per-structure UMA and DFT hull distances (eV/atom) for the 88 spot-checked structures. Tier T is trusted and S is sparse by the UMA hull. Refs is the number of reference phases within 1 meV/atom of the DFT hull in the structure’s chemical subsystem. Agree (y or n) marks metastability claims on the same side of 0.1 eV/atom.

setting	value
groups per iteration $B$ , group size $K$	4, 16
time grid, species noise $\eta$	64 points, 0
rollout iterations, inner epochs	750, 3
optimizer, learning rate, gradient-norm clip	AdamW, $10^{-4}$ , 0.5
clip range $\epsilon_{\mathrm{clip}}$ of Eq. 3	0.2
channel weights $\alpha_{\mathrm{species}},\ \alpha_{\mathrm{pos}},\ \alpha_{\mathrm{cell}}$	1, 0.1, 0.1

Appendix

Table 12: Hyperparameters of the OMatGRPO run. The discovery run of Section 4.3 uses the same values with the single-element guard removed.

Explore similar work

CrystalGRPO: Target-Aligned and Coverage-Preserving Reinforcement Learning for Flow-Based Crystal Structure Prediction
Aug 6, 2026Kaixiang Su, Hongfei Xue, Qiang ZhuCrystal Structure Prediction Group Relative Policy Optimization
Department of Computer Science, University of North Carolina at Charlotte · Department of Mechanical Engineering and Engineering Science, University of North Carolina at Charlotte · North Carolina Battery Complexity, Autonomous Vehicle and Electrification (BATT CAVE) Research Center
Riemannian Flow Models with Reinforcement Learning for Molecular Crystal Structure Prediction
Sep 30, 2026Thomas Egg, Harry Winston Sullivan, Maya M. Martirossyan +8Crystal Structure Prediction Crystal Structure Generation
Center for Soft Matter Research, Department of Physics, New York University, New York 10003, USA · Simons Center for Computational Physical Chemistry, Department of Chemistry, New York University, New York 10003, USA · Department of Chemical Engineering and Materials Science, University of Minnesota, Minneapolis, MN 55455, USA +7
Topology-Stratified Materials Discovery with A Flow-Based Generative Model
Sep 22, 2026Jingyi Zhou, Oyshee Chowdhury, Noah Oyeniran +1Materials Science Generative Modeling
Department of Aerospace Engineering and Mechanics, University of Alabama, Tuscaloosa, 35487, Alabama, United States. · Alabama Materials Institute, The University of Alabama, Tuscaloosa, 35487, Alabama, United States.

Aug 6, 2026J/K move · Enter open · S save

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

Kaixiang Su, Hongfei Xue, Qiang Zhu

Department of Computer Science, University of North Carolina at Charlotte · Department of Mechanical Engineering and Engineering Science, University of North Carolina at Charlotte · North Carolina Battery Complexity, Autonomous Vehicle and Electrification (BATT CAVE) Research Center

Crystal Structure Prediction Group Relative Policy Optimization Generative Flow Networks RL Post-Training

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Riemannian Flow Models with Reinforcement Learning for Molecular Crystal Structure Prediction

Topology-Stratified Materials Discovery with A Flow-Based Generative Model