[challenge]: Reinforcement fine-tuning for inverse materials design

Track: AI-agents-for-research · generative materials
Difficulty: advanced (comfortable reading ML papers and modifying PyTorch/JAX)
Compute: one GPU per team is enough for the week

Overview

A pretrained generative model of crystals already knows chemistry: give it a formula and it returns plausible structures. But plausible is not useful — out of the box it reproduces the average of its training set, with no notion that you want a crystal that frequency-doubles a UV laser, or one as light as possible without falling apart.

This is where GPT was before RLHF, and the fix is the same: reinforcement-fine-tune the generator against a reward that encodes what you want. Your task this week is to do that for crystals, using three ingredients:

1. a generative model — CrystalFormer (the policy)
2. a reward model — a property or stability predictor
3. reinforcement learning — PPO/GRPO to close the loop

Pick one of two starter hunts, or bring your own target:

• Track A — deep-UV nonlinear-optical crystals
• Track B — ultralight stable solids
• Track C — your own objective

The method: RL fine-tuning, as in LLM alignment

Same recipe as aligning an LLM — sample, score, reinforce — with a KL leash to the base model so the policy stays fluent.

	Aligning an LLM	This challenge
Policy	pretrained LLM	pretrained CrystalFormer (space group → Wyckoff → atoms → coords → lattice)
Data	internet text	~1.4 M known crystals (Alexandria / Materials Project)
Reward	learned (human preferences) or exact (code/math checker)	learned (property predictor) or exact (symmetry, density)
Algorithm	PPO / GRPO	PPO / GRPO (in the repo)
Target	helpful, harmless	stable and on-target property
Failure mode	reward hacking	reward hacking

The loop, from Cao & Wang (arXiv:2504.02367, Fig. 1a) — sample from CrystalFormer, score with a reward model, update by RL, repeat:

The objective is the paper's:

ℒ = 𝔼[ r(x) − τ · ln( p_θ(x) / p_base(x) ) ]

Maximize reward without drifting far from the pretrained model; the τ term is the KL leash.

A reward can be learned (a property predictor — flexible but hackable) or exact (symmetry, density — cheap and unfakeable). Both starter tracks are anchored by something exact — symmetry in Track A (enforced at generation), density in Track B — so a hit can be verified without expensive DFT.

Background

• CrystalFormer (Cao & Wang, IOP CAS, with the Lü Jian group at Jilin University; Science Bulletin 2025 — news, code) generates in the symmetry-reduced Wyckoff representation, so the space group is an explicit token you can steer. That is why the symmetry-based targets below fit so naturally.
• Reinforcement Fine-Tuning for Materials Design (Cao & Wang, arXiv:2504.02367) is the direct predecessor: property models as rewards, joint optimization of competing properties (e.g. dielectric constant and band gap). Read it first.
• The repo already ships PPO/DPO/REINFORCE, an end-to-end sampling → relaxation → stability pipeline, and reward factories. You extend a working system, not build from scratch.

Tracks

Tracks A and B are worked examples, ready to run; Track C invites your own objective. Comparing tracks, or combining properties, makes for a strong submission.

Track A — deep-UV nonlinear-optical crystals

NLO crystals frequency-double lasers; deep-UV ones (e.g. the borate KBBF) are rare and strategically important. Second-harmonic generation is forbidden in any centrosymmetric crystal, so the core requirement is a non-centrosymmetric space group. CrystalFormer emits the space group as its first token, so you get this by construction — condition generation on non-centrosymmetric groups rather than filtering after the fact. Stability and band gap then drive the reward, and spglib confirms the symmetry survives relaxation:

Ingredient	What it enforces	How	Role
non-centrosymmetric space group	SHG is symmetry-allowed	condition CrystalFormer's space-group token; verify with spglib after relaxation	constraint (by construction)
stability (E-hull)	the crystal could exist	MLFF relax + energy-above-hull (in repo)	reward — exact (MLFF)
wide band gap	UV transparency	a band-gap predictor (yours to provide)	reward — learned
large SHG coefficient (stretch)	it is actually a good NLO crystal	DFT on the top candidates	verify-only

Reward = stability × gap-in-window, over structures generated directly in non-centrosymmetric space groups. Fine-tune to raise the yield of stable, wide-gap candidates, then validate the best.

Track B — ultralight stable solids

What is the lightest thermodynamically stable crystal? Density is exact (mass / cell volume) and stability is the MLFF E-hull you already have — no learned model needed, so the reward is fully verifiable.

Reward term	What it enforces	How it's computed	Reward type
low mass density	lightness	mass / cell volume	exact
stability (E-hull)	it won't collapse	MLFF relax + energy-above-hull (in repo)	exact (MLFF)

Reward = f(1/density) with a stability penalty. The interest is the tension: low density favors open frameworks, stability punishes them. Map and push the density-vs-stability Pareto front.

Track C — your own objective

The recipe is general — thermoelectrics, transparent conductors, high-refractive-index optics, auxetics, ferroelectrics, and more are all fair game. A good self-chosen target meets three bars:

• useful — a real application to point at;
• cheap to score — seconds per candidate (fast predictor, empirical formula, or exact check), since RL needs thousands of evaluations and per-sample DFT is too slow as a reward;
• verifiable — a crisp way to confirm a hit (exact if possible, else a higher-fidelity model or DFT on the top candidates).

Clear those bars and go; check with the organizers early.

Getting started

Works end-to-end today:

• Pipeline: postprocess.sh runs sample → convert → MLFF-relax → energy-above-hull (ORB-v3 by default). This is your evaluator.
• RL: crystalformer/cli/train_ppo.py drives PPO with rewards (--reward {ehull,prop,dielectric,force}), exploration knobs (--sg_temperature, --sg_epsilon, --exploration_weight, --diversity_weight), a KL penalty (--beta), and entropy (--alpha). DPO and classifier-guided generation are also available.
• Rewards: crystalformer/reinforce/reward.py has make_ehull_reward_fn (stability), make_prop_reward_fn (scalar → target), and make_dielectric_reward_fn (dielectric × gap). Use the dielectric one as a worked reference for wiring in a property model, and make_prop_reward_fn as the scaffold for your own.

Your build:

• Track A: generate directly in non-centrosymmetric space groups — fix one with --spacegroup, or mask CrystalFormer's space-group distribution to the non-centrosymmetric subset — and reward stability × band gap. The band-gap predictor is yours to provide (a pretrained MEGNet/MatGL/matminer model, or a light one trained on Materials Project gaps). Use spglib post-relaxation to confirm the symmetry holds.
• Track B: add a density term — no external model needed; the work is reward shaping and the stability trade-off.

The reward model is part of the challenge, not a given: it is what separates a real discovery from a hallucinated one.

Deliverables

Bounded for a week, with a tail that can become a paper. A submission can be a research study, an agentic tool, or both.

Research study — the RL fine-tuning experiment:

1. A reproducible RL fine-tuning run: documented reward (terms, weights, rationale), training config, reward curves, and a sanity check that the reward scores known materials correctly.
2. A before/after result: sample from the base and tuned models; report the shift in the verifiable metric with error bars over several seeds.
   • Track A: fraction of samples that are non-centrosymmetric, stable, and in the gap window.
   • Track B: density distribution and lowest density at fixed stability.
3. Validated candidates: exact check on all samples; DFT relaxation + property on the top finalists; a table of the best (formula, space group, E-hull, property), marking rediscoveries versus novel.
4. A short report/notebook and a pull request, including one honest reward-hacking story — how the policy cheated, how you caught it, what you did.

Agentic framework — a human-in-the-loop design tool for experimentalists:

• takes a target stated in plain language and turns it into a reward/constraint spec;
• orchestrates the loop — generation, scoring, optional RL fine-tuning, validation — and returns candidate structures with provenance, stability and property estimates, and the reasoning behind each pick;
• builds on the repo's MCP server (mcp/csp_mcp.py, which exposes generate_structures): wrap the pipeline steps as agent tools behind a clean interface.

Deliver the working tool, a short recorded design session, and notes on reliability — above all, how it avoids handing an experimentalist hallucinated or unstable candidates.

Judging

• Correctness and reproducibility of the RL loop.
• Rigor of the before/after (real baselines and error bars).
• Depth of validation.
• Insight into reward hacking and the exploration–exploitation trade-off.
• Clarity.

Bonus: rediscovering a known target material, a validated novel candidate, or a clean multi-property Pareto analysis.

Reward hacking

The part that makes it research. The policy will find the cracks in your reward, as LLMs learn to flatter a flawed preference model: structures scored high-gap because they fall outside the predictor's training distribution; collapse to a single high-reward template; "stable-looking" structures the force field rejects. Detecting and defending — held-out DFT checks, reward-model ensembles or uncertainty penalties, diversity metrics, tuning the KL leash — is the science. This is why both starter tracks anchor on something exact — symmetry (enforced at generation in Track A) and density (Track B) cannot be hallucinated.

Research extensions

• Better RL: GRPO, or risk-sensitive GRPO — does risk-seeking surface rarer, better candidates?
• Multi-objective: a real Pareto front (A: symmetry × gap × SHG; B: density × stability × a mechanical constraint), as with dielectric × gap in the paper.
• More symmetry-defined targets: ferroelectrics/piezoelectrics (polar groups), chiral crystals (Sohncke groups) — all exact to check with spglib.
• Robust rewards: ensembles, uncertainty penalties, active learning from DFT feedback.
• Generalization: fine-tune on one chemical system, test transfer to unseen elements.

Resources

• CrystalFormer: https://github.com/deepmodeling/CrystalFormer
• RL fine-tuning paper (read first): https://arxiv.org/abs/2504.02367
• CrystalFormer background: https://www.iop.cas.cn/xwzx/kydt/202510/t20251022_7994908.html
• ML force fields: ORB / MACE / MatGL (train_ppo.py --mlff_model)
• Symmetry and analysis: spglib, pymatgen, matminer
• Concepts: RL fine-tuning of LLMs (RLHF / RLVR); PPO/GRPO; energy-above-hull; second-harmonic generation and non-centrosymmetric point groups

Read the cited reinforcement-fine-tuning paper first, then run postprocess.sh and inspect crystalformer/cli/train_ppo.py and crystalformer/reinforce/reward.py. Choose Track A or B and use the existing reward and sampling entry points as the starting map. Done means a reproducible run with documented configuration, before/after metrics, validation results, reward-hacking analysis, and a report or pull request.