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

1 INTRODUCTION

• 45.5% mSUN after reinforcement versus 13.4% for the pretrained model; best-of-N reaches 7.7%, and frozen-composition reinforcement reaches 9.8%.DFT agrees with the machine-learned potential on 63 of 67 spot-checked structures; the reward also raises Chemeleon2’s mSUN from 36.6% to 43.7% in its pipeline.Source: Figure 1 · p. 2, §1 · p. 3
• Directly reinforcing atom-type transition likelihoods addresses a gap because elemental constituents determine chemistry, yet prior approaches froze composition, relaxed it to continuous variables, or decoded it from latent space.The authors describe whether direct reinforcement of compositional decisions improves generation as an open question.Source: §1 · p. 2
• OMatG combines stochastic interpolants for positions and lattice with masked discrete flow matching for atom types, whose categorical transitions provide exact likelihoods for GRPO.Each atom type is chosen in a single categorical draw during inference, and the policy gradient acts directly on the composition channel.Source: §1 · p. 2
• Without guards, reinforcement can exploit the benchmark by repacking known single elements into new cells, which count as unique, novel, and metastable despite yielding no new compounds.Such structures made up 53.5% of the unguarded run’s mSUN set; the reward adds guards for single-element structures and sparse reference hulls.Source: §1 · p. 3
• The work combines discrete-channel GRPO, a guarded reward that transfers to another generator, and an analysis showing that single-element structures and sparse reference hulls can inflate mSUN.It also reports results by the number of reference phases behind each generated structure.Source: §1 · p. 3

2 RELATED WORK

• OMatG-IRL reinforced stochastic interpolants for positions and lattice while keeping species generation frozen; earlier crystal RL also acted on formulas, templates, or text representations.The cited work SPARC uses a reward-weighted denoising loss for discrete atom types but lacks a likelihood ratio, group baseline, or clipping.Source: §2 · pp. 3–4
• In recent three-dimensional diffusion and flow generators, composition was frozen as a condition, relaxed to continuous variables, or decoded from a latent rather than entering the policy gradient as a discrete variable.SPARC applies reward-weighted denoising to discrete atom types without a likelihood ratio, group baseline, or clipping.Source: §2 · p. 4
• Prior work describes reward hacking in reinforcement learning and diffusion fine-tuning, while crystal-generation studies report composition collapse and benchmark concerns including rediscovery, duplicates, and potential-dependent stability labels.The S.U.N. metric is standardized by LeMat-GenBench, and cited studies examine enumeration baselines and substitution variants of training structures.Source: §2 · p. 4

3.1 BACKGROUND: STOCHASTIC INTERPOLANTS AND MASKED DISCRETE FLOW MATCHING

• Stochastic interpolants transport positions and lattice from a base distribution toward crystal structures, while masked discrete flow matching progressively unmasks atom types using learned denoising probabilities.The atom-type process is a continuous-time Markov chain integrated on the same Euler time grid as the continuous SDE.Source: §3.1 · p. 4

3.2 GRPO ON THE NATIVE COMPOSITION CHANNEL

• GRPO compares groups of K rollouts from the same initial condition, assigning each rollout a final-structure reward and a group-relative advantage.The sampler’s transition kernel factors by positions, lattice, and species; channel weights compensate for different update scales, with a KL penalty toward the frozen pretrained model.Source: §3.2 · pp. 4–5
• For composition, the one-step policy ratio is exact: schedule-only factors cancel, leaving the predicted element probabilities for atoms that commit at that step.The product’s factors are update points for the policy gradient, which therefore reaches the composition channel exactly.Source: §3.2 · p. 5

3.3 REWARD STACK AND GUARDS

• Every rollout structure is relaxed before scoring; the stability term uses energy above the UMA-consistent convex hull, abstaining when fewer than 12 near-hull references support the chemical subsystem.Single-element structures and structures more than 0.1 eV/atom below the hull receive the penalty floor.Source: §3.3 · pp. 5–6
• Table 1 reports mSUN over the nominal 2,500 generations, while unique, novel, and metastable rates use valid structures and SUN and single-element mSUN are counts.The caption notes that best-of-N’s zero single-element and sparse-hull entries are enforced by its selection score rather than learned.Source: Table 1 · p. 6
• The reward credits novelty and uniqueness, discounts repeated compositions within rollout groups, and uses a batch-level MMD bonus to discourage chemistry-wide mode collapse.Novelty is checked against the MP-20 training set, while uniqueness is checked against earlier same-formula structures in the rollout batch.Source: §3.3 · p. 6

4.1 SETUP

• Runs start from an MP-20-pretrained OMatG model and train for 750 rollout iterations, with each batch reused for three inner policy updates.The benchmark defines mSUN as the fraction of generated structures that are metastable, unique, and novel at once.Source: §4.1 · p. 6

4.2 RL TRIPLES THE YIELD OF STABLE, NEW COMPOUNDS

• 45.5% mSUN (1,138 of 2,500), up from the prior’s 13.4% (336 of 2,500), a 3.4× gain; Chemeleon2 reaches 36.6%.Every method submits 2,500 structures for LeMat-GenBench scoring; Figure 2a separates trusted from sparse hulls.Source: Figure 2 · p. 7, §4.2 · p. 7
• Best-of-N reaches 7.7% mSUN and frozen-composition reaches 9.8%, while the controls suggest most of the full gain comes from species reinforcement.Across the trusted-tier spot-checks, DFT confirms 63 of 67 metastability claims, with a 0.016 eV/atom mean absolute deviation between UMA and DFT hull distances.Source: §4.2 · p. 7

4.3 REWARD HACKING

• Clipping the stability reward removed the below-hull incentive, but withholding reward on sparse hulls led the policy toward those regions and single-element structures.In the discovery run, sparse-hull generations rose from 0.19 to 0.58, and 53.5% of its mSUN set was single-element.Source: §4.3 · p. 8
• The single-element penalty guard reduces single-element structures to 0.4% of OMatGRPO’s mSUN set and keeps their training fraction near zero.The sparse-routed fraction remains near 0.2 during training, and the guard holds across three seeds.Source: §4.3 · p. 8
• The reward and benchmark share components, but without the creativity term the reward still reaches 39.3% mSUN versus the prior’s 13.4%.The reward scores novelty against MP-20, whereas the benchmark tests it against roughly five million LeMat-Bulk structures; DFT is the one stability check unused by the reward.Source: §4.3 · p. 8

4.4 SWAPPING REWARDS ACROSS GENERATORS

• The proposed reward gives higher mSUN in both pipelines: 45.5% versus 18.6% in OMatGRPO and 43.7% versus 36.6% in Chemeleon2.Table 2 swaps rewards between the pipelines while keeping each host pipeline’s relaxation convention fixed.Source: Table 2 · p. 9, §4.4 · p. 9

4.5 A STABILITY CLAIM IS ONLY AS GOOD AS ITS REFERENCE HULL

• Sparse reference hulls limit what metastability claims establish: 63.4% of the discovery run’s mSUN set is on hulls with fewer than 12 references, versus 23.5% for OMatGRPO.On sparse hulls, DFT tracks UMA but cannot adjudicate stability because the DFT hull is missing the same phases; the paper reports sparse-tier mSUN separately.Source: §4.5 · p. 9

5 CONCLUSION

• LeMat-GenBench yield of metastable, unique, novel structures rises from 13.4% to 45.5%; DFT confirms 63/67 spot-checked stability claims, while the same reward raises Chemeleon2 from 36.6% to 43.7% in its pipeline.Source: §5 · p. 9
• Every guard closed an exploit, including repacking single elements into new cells, which mSUN implementations count as discovery despite yielding no new compounds.The authors recommend treating sparse-hull mSUN as candidates rather than discoveries.Source: §5 · pp. 9–10
• The authors argue benchmarks should report the number of reference phases behind each structure, since stability claims depend on the reference hull.They suggest computing missing competing phases with DFT to resolve promising sparse-hull candidates.Source: §5 · p. 10
• The policy objective, reward, benchmark protocol, DFT protocol, code, trained models, and generated structures are identified as reproducibility resources.The paper also records the reward-swap protocol and deviations.Source: REPRODUCIBILITY STATEMENT · p. 10
• The work uses public materials data and targets stability, novelty, and diversity of inorganic crystals, with no identified risks beyond computational materials discovery generally.The authors report results split by reference-phase counts so readers can judge the support for stability claims.Source: ETHICS STATEMENT · p. 10
• Claude assisted with implementation and review of code, experiments, figures, literature review, derivation drafting, and text; the authors reviewed the work and take responsibility for it.It was not used to generate data.Source: LLM USAGE DISCLOSURE · p. 10

A DERIVATION OF THE POLICY RATIO AND THE KL TERMS

• The appendix derives the composition-channel ratio of Eq. 4 and the KL terms of Eq. 3 for any species-noise value η ≥ 0.The derivation follows the continuous-time Markov-chain construction of masked discrete flow matching and the OMatG species model.Source: A DERIVATION OF THE POLICY RATIO AND… · p. 15
• A structure has N atoms, each with a species state that can be the mask or one of S species; generation starts fully masked and ends with no atoms masked.The sampler takes T Euler steps of size ∆t.Source: A DERIVATION OF THE POLICY RATIO AND… · p. 15
• The model prediction is the final species distribution for each atom, while the reference policy is the frozen pretrained checkpoint run with the same η and time grid.For the composition channel, positions and lattice have already advanced to the next step when the species state is evaluated.Source: A DERIVATION OF THE POLICY RATIO AND… · p. 15

A.1 MASKED DFM AS A CONTINUOUS-TIME MARKOV CHAIN

• The masked interpolant gives each atom its final species with probability t and otherwise leaves it masked; generation uses a Markov chain whose model-dependent transition is unmasking.The remask rate is sampler-fixed, and there is no direct transition between species.Source: §A.1 · pp. 15–16

A.2 THE EULER STEP KERNEL

• Each Euler step factorizes over atoms, drawing unmask and remask decisions as Bernoulli variables and choosing an unmasking atom’s species from the model prediction.All runs use η = 0, so an unmasked atom keeps its species with probability one; model predictions enter only for atoms that unmask in that step.Source: §A.2 · p. 16

A.3 THE POLICY RATIO

• The composition-channel ratio reduces to contributions from atoms that commit at a step; at η = 0 each atom commits once, while remasking at η > 0 creates additional update points.At the final integration step, any remaining masked atoms are forced to unmask and enter the commit set with weight one.Source: §A.3 · p. 16

A.4 KL TERM FOR THE COMPOSITION CHANNEL

• The composition-channel KL is a sum of per-atom divergences and vanishes for unmasked atoms; for masked atoms it is the unmask probability times the categorical KL between species predictions.This form holds for any η ≥ 0, and the unmask-probability factor reaches one at the last step.Source: §A.4 · p. 16
• The implementation evaluates this KL at every masked atom of every step, whether or not it commits, using η = 0 and β_species = 0.05.At η = 0, the unmask probability is min{1, ∆t/(1−t)}.Source: §A.4 · p. 17

A.5 CONTINUOUS CHANNELS

• Position steps have Gaussian kernels, and their KL term is based on squared drift differences; the lattice uses the same kernel form but has no KL term.The position variance is floored at 10^-4, and β_pos = 0.01 while β_cell = 0.Source: §A.5 · p. 17

Internal evaluation

• The internal evaluation generates 1,000 structures per checkpoint and counts those that are metastable, unique, and novel under its stated relaxation and matching rules.Metastability uses Ehull ≤ 0.1 eV/atom; novelty is measured against the MP-20 training split, or Alex-MP-20 for the Alexandria-pretrained model.Source: B INTERNAL EVALUATION AND CONVENTION MAP · p. 17

Convention map

• Table 3 places the internal and benchmark conventions side by side, but their counts are not interchangeable; every quoted internal count has a benchmark counterpart.The appendix says the structure-by-structure join was checked against the benchmark stability values without disagreement.Source: B INTERNAL EVALUATION AND CONVENTION MAP · p. 17, Table 3 · p. 18

Benchmark results of every named run

• Routing sparse-reference structures to the floor removes the single-element exploit, but yields less than the guard.The pre-creativity runs differ from their counterparts only by the creativity term.Source: B INTERNAL EVALUATION AND CONVENTION MAP · pp. 17–18

The creativity ablation

• Without creativity, the guarded reward reaches 39.3% mSUN versus the prior’s 13.4%; adding creativity raises it to 45.5%.Novelty was rewarded against the MP-20 training split but evaluated against LeMat-Bulk; uniqueness is duplicate suppression in both settings, not independent evidence.Source: B INTERNAL EVALUATION AND CONVENTION MAP · p. 18

Mode 1, displacement

• At strict force tolerance, the displacement reward’s within-group spread collapsed to 0.045 Å, against 1.63 eV/atom for energy-family rewards, and learning stopped.With loose relaxation, the policy instead produced near-stationary structures with implausible energetics.Source: C REWARD-HACKING CATALOG WITH PER-RUN… · p. 19

Mode 2, formation energy

• Formation-energy optimization lowered median formation energy to −2.31 eV/atom while producing the compared run’s worst median Ehull, +1.08 eV/atom.The fluorine-containing share rose from 14% to 91%, while formation energy and hull distance remained weakly correlated (Pearson +0.09 to +0.36).Source: C REWARD-HACKING CATALOG WITH PER-RUN… · p. 19

Mode 3, unfloored Ehull

• Unfloored Ehull optimization chased below-hull depth in Li–O compositions whose reference hulls contained too few competing phases.Flooring the term at zero makes on-hull and below-hull structures tie at the best value, so fictitious depth earns nothing.Source: C REWARD-HACKING CATALOG WITH PER-RUN… · p. 19

Mode 4, sparse above-hull binaries

• Routing structures with fewer than 12 near-hull references away from credit closed sparse-binary drift that persisted despite composition-channel KL anchoring.The held-out fluorine share fell from 12% to 7%, while the median reference count stayed near 32 rather than falling to 3.Source: C REWARD-HACKING CATALOG WITH PER-RUN… · p. 19

Mode 5, single-element packings

• Under internal evaluation, discovery runs produced 253 and 242 single-element mSUN structures per 1,000 generations, versus 6 for the prior; the guard reduced the pre-creativity count to 12.Across three guarded seeds the count ranged from 2 to 15, and OMatGRPO had 2.Source: C REWARD-HACKING CATALOG WITH PER-RUN… · pp. 19–20
• The Alexandria prior produced no single-element mSUN structures, but pre-creativity training induced 104 per 1,000; larger groups raised MP-20 counts from 253 at K=16 to 307 at K=32.The group-size comparison used the same rollout budget, and both comparisons preceded the creativity term.Source: C REWARD-HACKING CATALOG WITH PER-RUN… · p. 20, Figure 4 · p. 20

Cross-protocol confirmation

• MatterGen’s released evaluation classified 16 of the first 20 offending packings as stable, unique, and novel; single-element structures made up 53.5% of the discovery run’s mSUN set.Chemeleon2’s mSUN definition also accepts these packings.Source: C REWARD-HACKING CATALOG WITH PER-RUN… · p. 20

One collapse that is not gaming

• Without the coverage bonus, emitted-element entropy narrowed to roughly 1.7; with the bonus, it stayed at least 3 across all 750 rollouts.The paper distinguishes this collapse from an exploit.Source: C REWARD-HACKING CATALOG WITH PER-RUN… · p. 21

Reward-swap setup

• The swaps import Chemeleon2’s four reward components unchanged or install our full reward with its routing, while preserving each host’s released training recipe.Each run generated 2,500 structures and was scored with the pinned LeMat-GenBench install; all runs were single-seed.Source: D REWARD-SWAP DETAILS · p. 21

Results

• Adding relaxation raises Chemeleon2’s reward in our pipeline from 0.160 to 0.186, but lowers our reward in their pipeline from 0.437 to 0.235.The comparison holds each host pipeline’s relaxation convention fixed so a reward swap changes only the reward.Source: D REWARD-SWAP DETAILS · p. 21

Interpretation of the two off-diagonal cells

• A′ shifted its mSUN set toward quaternary and quinary chemistry, had a 2.1% sparse-hull share, and reached the study’s highest strict-SUN count of 52.The same reward did not produce these traits in the authors’ pipeline; the authors say replay diagnostics and a seed replicate would test and bound their interpretation.Source: D REWARD-SWAP DETAILS · pp. 21–22
• Only about 30% of Chemeleon2’s within-group learnable signal sits in its energy term, while most gradient pushes novelty and coverage rather than hull proximity.In the frozen-VAE pipeline, the decoder already supplies geometric sanity, so the reward’s steering budget goes to composition; the authors attribute the observed traits to a reward–decoder interaction, not the reward alone.Source: D REWARD-SWAP DETAILS · pp. 21–22

Did our pipeline hack the transplanted reward?

• The B run’s reward rose about 17% mainly through composition diversity, while metastability moved from 0.272 to 0.321; its mSUN set was 7.5% single-element structures.Creativity was already at its ceiling, contributing no within-group variance; relaxation before scoring reduced the single-element share to 4.1%.Source: D REWARD-SWAP DETAILS · p. 22, Table 9 · p. 22

• The 88-structure spot-check used three predeclared batches and VASP 6.4.1 with PBE relaxation and static calculations against the released DFT LeMat-Bulk hull.Two Gd-bearing calculations were excluded and not repeated; neither entered any statistic.Source: E DFT PROTOCOL AND SPOT-CHECK TABLES · p. 23, Table 11 · p. 24
• Four trusted-tier and two sparse-tier metastability claims flip between UMA and DFT; sparse reference hulls and threshold-adjacent values qualify the agreement.Examples include Mg2Yb4 changing from −0.012 to +0.155 eV/atom and Rb6 from 0.061 to 0.186 eV/atom.Source: E DFT PROTOCOL AND SPOT-CHECK TABLES · p. 23, Table 11 · p. 24
• UMA–MatterSim differs by a median 0.015 eV/atom after a composition-linear offset on 990 of 1,000 mode-4 development-run structures, and was not repeated for OMatGRPO.The comparison used UMA-relaxed structures and removed elemental reference-energy offsets by least squares.Source: E DFT PROTOCOL AND SPOT-CHECK TABLES · p. 23
• Trusted-tier claims agree 63/67 at the 0.1 eV/atom threshold (MAE 0.016 eV/atom); sparse-tier claims agree 19/21 (MAE 0.020 eV/atom), despite shared missing phases.The sparse-tier agreement validates energies, not the stability of those compositions.Source: E DFT PROTOCOL AND SPOT-CHECK TABLES · p. 23, Table 10 · p. 24

Model, sampler and optimization

• The OMatG run uses 750 rollout iterations, each with three optimizer epochs, and the discovery run uses the same hyperparameters with the single-element guard removed.Each iteration draws four groups of 16 structures; training uses AdamW with learning rate 10^-4 and gradient-norm clipping at 0.5.Source: F TRAINING CONFIGURATION AND BEST-OF-N… · p. 25, Table 12 · p. 25

Relaxation before scoring

• Relaxing every rollout structure before scoring is necessary because raw geometries shift by about 1 Å RMSD and yield far fewer known-structure matches.At most 0.2% of raw geometries match a known structure, versus 4–13% after relaxation; structures receive up to 100 FIRE steps under UMA.Source: F TRAINING CONFIGURATION AND BEST-OF-N… · p. 25

Stability term and routing

• Structures with fewer than 12 reference phases are routed to abstention, while single-element, excessively below-hull, or failed-relaxation structures receive the stability floor.The hull is the UMA-consistent LeMat-Bulk hull; sparse structures retain reward values but receive zero advantage and are excluded from group statistics.Source: F TRAINING CONFIGURATION AND BEST-OF-N… · pp. 25–26

Creativity term

• Creativity rewards structures that are both novel against the MP-20 training split and unique within the rollout batch, with distance-based partial credit otherwise.The partial credit is the minimum average-minimum-distance to matching structures, for k = 100, clamped to [0, 1]; a matcher call exceeding three seconds scores zero.Source: F TRAINING CONFIGURATION AND BEST-OF-N… · p. 26

Diversity terms

• The occurrence discount pulls repeated compositions toward the reward floor, while the batch-level MMD bonus rewards compositions that increase discrepancy from MP-20 compositions.Occurrence discount begins falling above three copies and reaches zero at six; MMD credits are rescaled to [0, 1] and weighted by 0.4.Source: F TRAINING CONFIGURATION AND BEST-OF-N… · p. 26

Frozen-composition control

• The frozen-composition control trains only position and lattice channels, sharing one pretrained-model composition across each group of 16 structures.It omits creativity, occurrence, and MMD terms; the stability pathway retains the floor, cap, sparse abstention, and relaxation limit.Source: F TRAINING CONFIGURATION AND BEST-OF-N… · p. 26

Best-of-N definition

• Best-of-N ranks 47,795 pretrained samples by the training reward’s stability pathway and submits the top 2,500 after cell-inclusive relaxation.Sparse-hull and single-element samples are excluded by the scoring-time floor routing; the same ranking cut at 1,000 yields a 4.4% mSUN rate.Source: F TRAINING CONFIGURATION AND BEST-OF-N… · p. 26