Reinforcement Learning for Crystal Relaxation

By jmedina9 • 12 min read • RL Environments

TL;DR

This project trains language models with reinforcement learning to generate relaxed crystal structures directly as CIF files. A fine-tuned CHGNet model provides formation-energy and force signals, while separate rubric and heuristic checks handle validity, composition, and bond geometry, making crystal generation a verifiable RL task. The main result is that model scale changes not just performance, but reasoning style: the larger model develops richer crystallographic reasoning, while the smaller model compresses its reasoning and plateaus on the hardest physics objective.

This project trains LLMs with reinforcement learning to directly propose relaxed crystal structures in a single forward pass, trying to either approximate DFT results or complement it. Crystal relaxation is a verifiable task: generating a stable crystal structure is difficult, but checking one is computationally cheap using a neural network potential. This verification asymmetry makes the problem a natural fit for Group Relative Policy Optimization (GRPO), which requires only a scalar reward from a verifier and no learned value function. A fine-tuned CHGNet model provides the formation-energy and force signals inside that verifier, while separate rubric and heuristic checks handle the rest, and the resulting analysis at two model scales reveals that size of the model determines the qualitative reasoning strategy rather than simply modulating performance magnitude.

The Challenge: Stable Structures Occupy a Tiny Corner of Structure Space

Reliable generation of stable crystal structures is a key enabler for closed-loop materials discovery. Physically plausible candidates occupy a small subset of all periodic atomic arrangements, near local minima of the quantum-mechanical energy landscape, while structure search and first-principles validation are computationally expensive. An alternative that is gaining a lot of interest is the use of generative models, generative models aim to learn the distribution of known stable crystals and rapidly propose new candidates that can subsequently be verified and relaxed.

A variety of deep generative approaches have been explored for crystalline materials, including variational autoencoders (VAEs), generative adversarial networks (GANs), and more recent diffusion-based models designed to respect symmetry and periodicity. Ensuring validity and stability remains challenging in part because crystal representations must handle rotational and translational invariance, permutation invariance, and periodic boundary conditions.

Among symmetry-aware diffusion approaches, the Crystal Diffusion Variational Autoencoder (CDVAE) [1] formulates crystal generation as a conditional score network that generates 3D atomic structures through a diffusion and denoising process. In practice, samples from CDVAE are produced by annealed Langevin dynamics that iteratively denoise noisy atom coordinates conditioned on a latent code, effectively acting as a learned refinement step that pushes structures toward lower-energy configurations. In downstream workflows, generated structures are typically further relaxed and validated with DFT to assess stability.

With the appearance of reasoning language models and the recent push for using transformers in crystal generation, this project asks: can a large language model trained with reinforcement learning learn to propose relaxed crystal structures directly? I reformulate the problem as a verifiable task: a setting where generating a correct solution is hard, but checking one is straightforward. Given a proposed CIF string^[7], a neural network potential can evaluate its formation energy and forces, and bond validity can be checked by fast heuristics, providing a deterministic reward signal without any learned critic. This verification asymmetry is the same property that has made RL effective for mathematical reasoning [4] and code generation, where unit tests serve as verifiers. Here, the verifier combines a fine-tuned CHGNet model for formation energy and forces with separate rubric checks, and the RL algorithm is GRPO, which samples multiple candidate structures, scores each against the verifier, and updates the policy using relative rankings within the group.

Environment Highlights:

State Space: CIF string encoding lattice vectors (a, b, c, α, β, γ), space group, and fractional Wyckoff coordinates for each atom. Dimensionality scales with system size (typically 3N + 6 degrees of freedom).
Action Space: Discrete token space over CIF text. Effectively encodes continuous atomic displacements (δx, δy, δz per atom) and lattice strain tensors (ε_ij) through text generation.
Reward: Relative-weighted sum R ∝ 0.10·r_fmt + 0.15·r_comp + 0.60·r_{E_f} + 0.24·r_F + 0.15·r_bond, with those weights renormalized internally before aggregation by a CHGNet inference server wrapped in the Prime Intellect verifiers library.
Training algorithm: GRPO samples G rollouts per prompt and optimizes the policy relative to the group mean reward, without a learned critic.

The Dataset

The training environment uses a curated subset of the Materials Project mp-20 benchmark [3], originally compiled by Chen et al. (2021). The dataset is downloaded directly from HuggingFace: 4everStudent/mp-20-relaxation-dataset-with-clusters. MP-20 contains ~45,000 experimentally reported inorganic crystal structures with up to 20 atoms per unit cell, spanning 89 elements across the periodic table.

Filtering criteria: To restrict the agent's search space to thermodynamically accessible materials, I applied a hard threshold of E_f < 0.1 eV/atom on the DFT formation energies from the Materials Project. This retains structures that lie within the commonly accepted convex hull tolerance for synthesizability [4], discarding high-energy metastable phases that would produce noisy reward signals during early training.

Dataset Statistics (filtered subset)

Total structures: ~45,000 (full mp-20) → ~28,000 after E_f filter
Training pool used: 2,000 randomly sampled structures
Average atoms/cell: 8.3
Space groups covered: 148 (out of 230 possible)
Formation energy range: −3.9 to 0.1 eV/atom
Composition diversity: unary through quaternary systems

🔮 The Model: Fine-tuned CHGNet

At the heart of this environment is CHGNet (Crystal Hamiltonian Graph Neural Network) [1], a universal MLIP pretrained on ~1.5 million Materials Project DFT calculations. CHGNet is a message-passing GNN that operates on crystal graphs where nodes are atoms and edges encode pairwise interactions within a cutoff radius (default 5 Å). A key architectural feature is the explicit inclusion of magnetic moments as node features alongside atomic species and fractional coordinates—enabling CHGNet to distinguish spin states that conventional force fields miss.

The model predicts four targets simultaneously: total energy E, per-atom forces F = −∂E/∂r, virial stress tensor σ, and site magnetic moments m. This multi-task formulation is critical for RL reward computation: I use E for the formation energy rubric and ‖F‖ for the forces rubric.

To reduce systematic bias on the filtered mp-20 subset, I fine-tuned the pretrained CHGNet checkpoint:

Fine-tuning Configuration:

Training samples: 2000 structures from the filtered mp-20 subset
Targets: Formation energy per atom + atomic forces
Batch size: 8
Epochs: 10
Optimizer: Adam with learning rate 1e-2
Loss function: Mean Squared Error (MSE)
Scheduler: Cosine annealing learning rate
Train/Val/Test split: 90% / 5% / 5%

Fine-tuning reduces the systematic offset (bias) that the pretrained model accumulates when evaluated on the low-energy filtered subset—the baseline CHGNet was trained on the full Materials Project including high-energy structures, so its prediction distribution is shifted. After fine-tuning, the MAE on formation energy dropped from ~0.21 eV/atom (baseline) to ~0.08 eV/atom on the test split, keeping reward noise within the 0.1 eV/atom threshold that defines my stability criterion.

🔬 Curious about the optimal hyperparameters?

I tested 72 different configurations to find the best setup. View the complete hyperparameter sweep results →

Below is the complete fine-tuning pipeline showing how CHGNet is trained on the MP20 subset.

1. Configuration & Imports

import numpy as np
from datasets import Dataset
from pymatgen.core import Structure
from pymatgen.ext.matproj import MPRester
from chgnet.data.dataset import StructureData, get_train_val_test_loader
from chgnet.model import CHGNet
from chgnet.trainer import Trainer

# Configuration
MP_API_KEY = os.getenv("MP_API_KEY")
DATASET_PATH = "mp_20.json"
N_TRAIN_SAMPLES = 2000 # Real challenge (use 500 for medium)
BATCH_SIZE = 8
EPOCHS = 10 # Real challenge (use 5 for medium)
LEARNING_RATE = 1e-2

2. Data Preparation

# Load dataset and sample
dataset = Dataset.from_json(DATASET_PATH)
subset = dataset.select(np.random.choice(len(dataset), N_TRAIN_SAMPLES, replace=False))

structures, energies_per_atom, forces = [], [], []
base_model = CHGNet.load()  # Load pretrained model

# Fetch structures and prepare training data
for data in tqdm(subset, desc="Preparing data"):
    structure = get_structure_from_mp_id(data["id"])
    if isinstance(structure, Structure):
        structures.append(structure)
        energies_per_atom.append(data["formation_energy_per_atom"])
        forces.append(base_model.predict_structure(structure)["f"])

3. Create DataLoaders

chgnet_dataset = StructureData(
    structures=structures,
    energies=energies_per_atom,
    forces=forces,
)
train_loader, val_loader, test_loader = get_train_val_test_loader(
    chgnet_dataset, 
    batch_size=BATCH_SIZE, 
    train_ratio=0.9, 
    val_ratio=0.05
)

4. Fine-tune the Model

trainer = Trainer(
    model=base_model,
    targets="ef",  # energy + forces
    optimizer="Adam",
    scheduler="CosLR",
    criterion="MSE",
    epochs=EPOCHS,
    learning_rate=LEARNING_RATE,
    use_device="cpu",
    print_freq=6,
)

trainer.train(train_loader, val_loader, test_loader)
fine_tuned_model = trainer.model

# Save the fine-tuned model
torch.save(trainer.best_model.state_dict(), "fine_tuned_chgnet_model.pth")

5. Evaluation

# Evaluate on test set
y_true, y_pred = [], []
for data in sample_eval:
    structure = get_structure_from_mp_id(data["id"])
    if isinstance(structure, Structure):
        y_true.append(data["formation_energy_per_atom"])
        y_pred.append(fine_tuned_model.predict_structure(structure)["e"])

# Compute metrics
r2 = r2_score(y_true, y_pred)
mae = np.mean(np.abs(np.array(y_true) - np.array(y_pred)))
print(f"R² score: {r2:.4f}, MAE: {mae:.4f} eV/atom")

💡 Note: This training pipeline takes ~15-20 minutes on CPU. For production use, consider GPU acceleration and larger training sets.

6. Results Visualization

I compare the model's predictions before and after fine-tuning to demonstrate the improvement:

Before Fine-tuning

Pretrained CHGNet baseline

After Fine-tuning

Fine-tuned on 2000 MP20 structures (Real Challenge)

Side-by-side comparison showing improved accuracy after fine-tuning on the filtered MP20 subset.

🔬 Complete Hyperparameter Sweep Results

Want to see the full story? I conducted a comprehensive grid search over 72 configurations, testing different dataset sizes (100-2000 samples), batch sizes, learning rates, and epoch counts.

View Detailed Results & Analysis →

Includes: baseline metrics, top configurations, sample efficiency analysis, and replication script

The RL Environment: Multi-Turn Crystal Relaxation

The crystal relaxation task is framed as a MultiTurnEnv—a multi-step RL environment built on the Prime Intellect verifiers library. At each turn t, the LLM policy receives the full conversation history (initial structure + all previous proposals and rubric feedback) as context and outputs a new CIF. This is analogous to a chain-of-thought policy where the "scratchpad" is the turn history.

Training uses GRPO (Group Relative Policy Optimization): for each prompt, G = 8 rollouts are sampled from the current policy π_θ, their rewards are normalized within the group (subtract mean, divide by std), and the policy is updated to increase the probability of above-average rollouts. No value network is needed—the group statistics serve as an implicit baseline, reducing variance without introducing a separate critic.

🎯 Objective

Given a disordered initial structure, the agent must generate a valid CIF file that:

Has valid crystallographic format
Matches the target composition
Achieves formation energy below 0.1 eV/atom
Minimizes atomic forces (well-relaxed)
Maintains physically reasonable bond lengths

How It Works

1️⃣

Initial State: The environment provides a disordered crystal structure from the MP20 dataset with high formation energy.

2️⃣

Agent Action: The LLM agent generates a response containing:

Reasoning about structural improvements (in <|think_start|>...<|think_end|>)
A proposed CIF structure (in <|answer_start|>...<|answer_end|>)

3️⃣

Evaluation: CHGNet evaluates formation energy and forces, while separate rubric and heuristic checks handle bond lengths and the non-physics constraints.

4️⃣

Feedback: The environment provides detailed feedback on what passed/failed, guiding the agent toward better structures in subsequent turns.

Crystal Relaxation Across Turns

The LLM iteratively refines atom positions and lattice parameters, lowering formation energy each turn.

→ →

⚠ Unrelaxed Structure

Turn 1 — Partially Relaxed

Turn 2 — Relaxed ✓

Reward Structure

The total reward is a weighted sum of individual rubric scores. For a detailed discussion of how these weights were chosen and the pitfalls encountered along the way, see Lessons from Designing Rewards for Crystal Generation.

Rubric	Weight	Criteria
Format	0.10	Valid CIF syntax
Composition	0.15	Matches target formula
Formation Energy	0.60	< 0.1 eV/atom
Forces	0.24	Low atomic forces
Bond Lengths	0.15	Physically reasonable

Note: Formation energy has the highest relative weight (60%) as the primary indicator of structural stability, and the displayed weights are renormalized internally before aggregation so the effective mixture sums to 1.0.

📄 Format Rubric (Weight: 0.10)

Ensures the agent produces syntactically valid CIF files and properly uses the answer tags.

What It Checks:

CIF string is present between <|answer_start|> and <|answer_end|> tags
CIF can be parsed by pymatgen's CifParser
Structure contains valid crystallographic data
No critical syntax errors in the CIF format

Scoring:

reward = 0.1 if cif_successfully_parsed else 0.0

Why it matters: Without a valid CIF, the structure cannot be evaluated by CHGNet. This rubric ensures the agent learns the proper format before focusing on energy optimization.

⚗️ Composition Rubric (Weight: 0.15)

Verifies that the generated structure has the correct chemical composition in the right ratios. This rubric serves two distinct roles: (1) it provides an easy gradient signal early in training, when the model is too weak to consistently pass the stability rubrics—preventing the policy from receiving zero reward on every step; and (2) it acts as an anti-reward-hacking mechanism, ensuring the agent cannot achieve high scores by memorizing a single simple crystal that passes energy/force thresholds regardless of the requested composition.

What It Checks:

Composition matches the target formula (e.g., NaCl, Li₂O)
Element ratios are correct (reduced composition)
No extra or missing elements

Implementation:

# Uses pymatgen's ElementComparator
cif_composition = structure.composition.reduced_composition
target_composition = Composition(target_formula).reduced_composition

composition_matches = ElementComparator().are_equal(
    cif_composition,
    target_composition,
)
reward = 0.15 if composition_matches else 0.0

Why it matters: The chemical composition fundamentally determines material properties—an agent that generates NaCl₂ instead of NaCl has failed at a basic chemistry level. Keeping weight low (0.15 vs. 0.60 for energy) prevents the agent from gaming composition alone, but non-zero weight ensures the rubric provides early training signal before the harder stability objectives become reachable.

⚡ Formation Energy Rubric (Weight: 0.60)

The most important rubric—evaluates thermodynamic stability through formation energy.

What It Checks:

Formation energy per atom < 0.1 eV/atom (threshold)
Calculated using fine-tuned CHGNet model
Lower energy = more stable structure = higher reward

Available Reward Strategies

The environment supports three different reward functions to guide the optimization process:

1. Discrete Reward (Step Function)

Binary feedback. Simple but provides sparse gradients.

def formation_energy_discrete(..., threshold=0.1):
    if energy <= threshold:
        return 1.0
    return 0.0

2. Linear Reward (Default)

Recommended

Continuous feedback. Decays linearly as energy rises above threshold.

def formation_energy_linear(..., threshold=0.1, max_penalty=2.0):
    if energy <= threshold:
        return 1.0
    elif energy >= threshold + max_penalty:
        return 0.0
    else:
        # Interpolate 1.0 -> 0.0
        return 1.0 - (energy - threshold) / max_penalty

3. Exponential Reward

Aggressive penalty for high energy structures.

def formation_energy_exponential(..., threshold=0.1, scale=1.0):
    if energy <= threshold:
        return 1.0
    excess = energy - threshold
    return np.exp(-excess / scale)

⚠️ Critical Threshold

Structures with formation energy above 0.1 eV/atom are considered thermodynamically unstable and unlikely to exist naturally. The agent must learn to minimize this value through proper atomic positioning and lattice parameters.

Why it matters: Formation energy is the gold standard for predicting whether a material can be synthesized. This 60% weight ensures the agent prioritizes structural stability above all else.

🔧 Forces Rubric (Weight: 0.24)

Measures how well-relaxed the structure is by examining forces on individual atoms.

What It Checks:

Maximum force on any atom < threshold (typically 0.5 eV/Å)
Average force magnitude across all atoms
Indicates if atoms are in equilibrium positions

Available Reward Strategies

The environment provides three ways to reward structural relaxation:

1. Discrete Reward

Strict binary check. Only rewards if the maximum force on any atom is below threshold.

def forces_discrete(..., threshold=0.05):
    # forces = get_forces(structure)
    max_force = np.max(np.linalg.norm(forces, axis=1))
    return 1.0 if max_force <= threshold else 0.0

2. Continuous Decay (Default)

Recommended

Smoothly penalizes larger forces. Even imperfect structures get some credit, guiding the gradient.

def forces_continuous_decay(..., decay=1.0):
    # forces = get_forces(structure)
    max_force = np.max(np.linalg.norm(forces, axis=1))
    return 1.0 / (1.0 + (decay * max_force))

3. Continuous Ratio

Rewards the percentage of well-behaved atoms. Good for large structures.

def forces_continuous_ratio(..., threshold=0.05):
    # forces = get_forces(structure)
    force_norms = np.linalg.norm(forces, axis=1)
    good_forces = np.sum(force_norms <= threshold)
    return good_forces / len(force_norms)

Physical intuition: In a stable structure, atoms sit at energy minima where forces are near zero. Large forces indicate atoms "want" to move—the structure hasn't truly relaxed.

💡 Agent learning: This rubric teaches the agent that low formation energy alone isn't enough—the structure must also be mechanically stable with minimal residual forces.

📏 Bond Lengths Rubric (Weight: 0.15)

Validates that interatomic distances are chemically reasonable based on atomic radii and bond types.

What It Checks:

Bond lengths fall within expected ranges (±32% tolerance)
Accounts for ionic vs covalent bonding (electronegativity difference)
Special handling for hydrogen bonds (more flexible)
Uses pymatgen's CrystalNN to identify bonded atoms

Available Reward Strategies

You can configure the strictness of the bond length check:

1. Discrete Reward

Binary accept/reject. Only rewards structures where the global bond score exceeds a threshold.

def bond_lengths_discrete(..., threshold=0.75):
    # score = bond_length_reasonableness_score(structure)
    return 1.0 if score >= threshold else 0.0

2. Continuous Reward (Default)

Recommended

Returns the raw reasonableness score (0.0 to 1.0), encouraging incremental improvements.

def bond_lengths_continuous(..., **kwargs):
    # score = bond_length_reasonableness_score(structure)
    if not score or score < 0:
        return 0.0
    # Returns raw score 0.0 - 1.0
    return max(0.0, min(1.0, score))

Common Failures:

Atoms too close → unphysical overlap
Atoms too far → weak/nonexistent bonding
Incorrect lattice parameters → compressed/expanded structure

Why it matters: Even if CHGNet predicts low energy, physically impossible bond lengths indicate the structure is chemically unrealistic. This rubric acts as a sanity check on the agent's geometric proposals.

🧪 Evaluation-Only Environment

In addition to the RL training environment described above, I created a separate evaluation benchmark designed specifically to measure LLM progress on crystal reasoning, materials science concepts, and structured scientific outputs.

This evaluation environment uses the same rubric system (format, composition, formation energy, forces, bond lengths) but applies it to a diverse set of tasks testing different aspects of crystal generation and reasoning—without any training involved.

Learn About the Evaluation Environment →

Includes: task design, evaluation philosophy, and use cases for LLM benchmarking

Run It on Prime Intellect

The environment is hosted on Prime Intellect and updated continuously. For the latest setup instructions, quickstart commands, and environment arguments, click the button at the top of this page.

View on Prime Intellect →

From MultiTurnEnv to RLM

Next phase

This post documents the original MultiTurnEnv version of crystal relaxation: chat-native feedback, rubric-guided retries, and answer extraction from the model response. The next iteration moves the task into an RLM harness where the agent works through files, checker calls, and a sandbox, with the final structure written to final.cif.

The benchmark becomes less sensitive to parser noise and more sensitive to actual structure quality.
Tool use, iteration strategy, and checker budgeting become part of the task itself.
Failures become easier to interpret: symmetry collapse, composition drift, force-proxy misses, and bad revision workflows.

Standalone follow-on page

I broke out the migration into its own page so the architecture change, eval results, and rollout examples can stand on their own.

Read the RLM transition →

Training Runs: How an LLM “Learns” Physics

• Tagged: RL Materials Science

I ran the crystal relaxation environment on two model scales: a 30B-parameter model (87 rollout samples, steps 0–100, 11 checkpoints) and a 230B-parameter model (63 samples, steps 0–070, 8 checkpoints). Both runs use the same reward function, the same filtered MP-20 dataset from HuggingFace, and the same GRPO training loop. What I found is that scale doesn’t just change how well a model learns—it determines what kind of reasoning strategy RL selects for entirely.

1. Reward Trajectories

Both models improve over training, but the trajectories look completely different:

Step	30B Model		230B Model
	Mean Reward	Std	Mean Reward	Std
0	0.406	0.202	0.706	0.364
10	0.710	0.235	0.705	0.145
20	0.703	0.214	0.880	0.164
30	0.498	0.272	0.722	0.270
40	0.727	0.223	0.502	0.308
50	0.696	0.276	0.719	0.186
60	0.589	0.275	0.954	0.116
70	0.529	0.267	0.928	0.144
80	0.766	0.277	—	—
90	0.661	0.327	—	—
100	0.520	0.258	—	—

The 230B model follows a non-monotonic but convergent trajectory: it dips at step 40 (mean 0.502) before recovering strongly to peak at step 60 (0.954) and sustaining near-peak performance at step 70 (0.928). Critically, its variance contracts—std drops from 0.364 at step 0 to 0.116 at step 60, meaning the model becomes more consistent as it learns. By step 60, even its worst sample (reward 0.692) would have been above-average at step 0.

The 30B model is more volatile. It peaks at step 80 (0.766) but regresses to 0.520 by step 100, with variance that never contracts below 0.20. Its best checkpoint is still 20% below the 230B model’s peak. The single most dramatic data point: at step 30, the 30B model produces the only zero-reward sample in either dataset (SnF₃, where the model correctly identifies a symmetry mismatch but fails to produce any valid output).

2. Scale Determines Reasoning Strategy

The most surprising finding comes from tracking the length and content of the models’ chain-of-thought reasoning blocks:

Step	30B reasoning (chars)	230B reasoning (chars)
0	1,532	3,966
20	1,502	4,440
40	1,270	5,969
60	924	6,640
70	810	7,320 (+85%)
80	654	—
100	574 (−63%)	—

The 230B model’s reasoning grows by 85% over training—from 3,966 to 7,320 characters by step 70. The 30B model’s reasoning shrinks by 63%—from 1,532 to 574 characters. These are not just different amounts; they are opposite directions. At step 80, one 30B sample produced a reasoning block of 3 characters—essentially skipping thinking entirely and going straight to CIF output. By step 100, the 30B’s maximum reasoning length (914 chars) is shorter than its step-0 minimum (1,288 chars).

System 1 vs. System 2 Reasoning

This maps cleanly onto Kahneman’s dual-process framework. The 230B model develops reasoning that is: deliberative, step-by-step crystallographic analysis that deepens over training. By step 30, it correctly identifies hexagonal cells from γ = 120°; by step 60–70, it references Wyckoff positions, prototype structures, and Pnma space groups it was never explicitly trained on—and explicitly acknowledges it cannot run real DFT while still producing stable structures.

The 30B model reasoning on the other hand is: compressed, procedural preambles (“relaxation will be performed using DFT—output below”) that discard domain vocabulary in favor of getting directly to the CIF. RL selects for this compression because the smaller model lacks the capacity to maintain both elaborate reasoning and correct structural output simultaneously. By step 90–100, the reasoning reads as a formality rather than a problem-solving step—the model even inserts the parenthetical “(simulated here for brevity)”.

Analyzing keyword data gives some evidence to this divergence. In the 230B model, “symmetry” mentions grow from 6.4 per sample at step 0 to 11.3 at step 60, and Wyckoff positions emerge at step 40+. In the 30B model, “symmetry” starts at 0.6 and reaches 0.0 by step 100. “Bond” drops from 2.0 to 0.0 by step 60. “Coordination” and “Wyckoff” never appear at any step.

3. Internalizing Bond Heuristics

The 30B model’s “bond” keyword count drops from 2.0 per sample at step 0 to 0.0 by step 60—and stays at zero through step 100. On the surface, this looks like the model not analyzing the structure entirely.

In reality, the bond length reward improves from 0.574 to 0.809 over the same period (peaking at 0.818 at step 90). The model’s bond quality gets better while it stops talking about bonds. This is the compression heuristic in action: the concept has been internalized into the structural output itself, removed from the scratchpad. The 30B model learns to do without needing to say. The 230B model does the opposite: bond mentions also decline after step 50, but from a higher baseline and after the model has already peak-loaded its symmetry vocabulary.

4. Conservative Perturbations—Two Very Different Kinds

Both models learn to make smaller lattice adjustments over time through step 60. But the degree, pattern, and outcome differ sharply:

Model	Δa at Step 0	Min Δa	Reduction	FE score at reward peak
30B	0.025 Å	0.002 Å (step 90)	−91%	0.658 (step 80)
230B	0.053 Å	0.017 Å (step 60)	−68%	0.957 (step 60)

The 230B model’s conservatism is healthy: it mirrors real DFT relaxation behavior, correlates with excellent formation energy scores, and notably rebounds at step 70 (Δa = 0.084 Å)—the model becomes more forceful once it has the structural reasoning to support larger corrections, and bond quality peaks at 0.938 that same step. The 30B model’s conservatism is pathological: modifying the lattice by 0.002 Å is essentially returning the input unchanged. It avoids catastrophic bond-length violations but cannot correct misplaced lattice parameters—hence formation energies that remain 10–20× above the 0.1 eV/atom threshold throughout all 100 steps.

5. Formation Energy: The Capacity Bottleneck

The formation energy reward component is where the capacity gap is most visible. The 230B model reaches a formation energy score of 0.957 at step 60 and sustains 0.896 at step 70. The 30B model peaks at 0.658 at step 80 and falls back to 0.283 by step 100—a ~45-percentage-point gap at peak, widening to 61 points by the end of each run.

In absolute terms, the 30B model’s generated structures have mean formation energies oscillating between 0.93 and 1.86 eV/atom across all steps—10–20× above the 0.1 eV/atom stability threshold. The structures pass composition and bond-length checks, but they are not genuinely relaxed. The model has learned to be a competent CIF formatter and a cautious structure perturber, but not a crystal-structure optimizer.

6. Effectiveness of Multi-Step Feedback

In multi-turn rollouts, when the first attempt fails the energy or the bond-length heuristic threshold, the model receives feedback and gets a second turn. I expected the model to use this feedback to diagnose and correct structural deficiencies. In practice, weak models show little evidence of their usefulness. Analysis of 37 first-turn failures from the 30B run:

Physical recoveries (energy, forces, bonds): 0. Zero cases where the model corrected formation energy, forces, or bond lengths.
Syntax recoveries (format corrections): 7 of 37. The model fixed a missing CIF line or formatting issue—it can repair a broken file, but not a bad structure.

The model treats scalar feedback (“Formation energy is too high”) as a signal to re-roll rather than to diagnose. It rephrases its chain-of-thought reasoning and outputs a near-identical CIF with minor random perturbations, having learned state recognition (what a valid crystal looks like) but not gradient estimation (how to move atoms toward lower energy).

On the positive side, inclusion of quantitative feedback about the formation energy improves the model’s predictive capabilities. Although it cannot fix a bad structure, when tasked to predict the formation energy, its performance improves significantly across turns—an “Actor-Critic” disconnect where the Critic (prediction) improves while the Actor (structure generation) stalls.

“The model learns to measure the problem before it learns to fix it.”

What this implies for the next run:

Scalar feedback (“energy too high”) provides no directional gradient in atom-coordinate space. The path forward is vector feedback—per-atom force vectors, specific bond-length violations, which Wyckoff site is strained—so the policy has an actual direction to move. The 230B model shows that with sufficient capacity, RL can induce genuinely calibrated structural reasoning. Whether a 30B model can exploit richer feedback to close the gap, or whether the formation energy bottleneck is a hard capacity limit, is the central open question.

What’s Next

Several directions are worth pursuing from here:

Benchmarks

A systematic benchmark comparing GRPO-trained LLMs against diffusion-based generative models (CDVAE, DiffCSP) and classical basin-hopping on the MP-20 test split, using formation energy, structural validity rate, and composition accuracy as evaluation axes.

Exploration of Reward Hacking Mechanisms

The composition-hacking failure mode I observed is one instance of a broader class. Systematic probing—varying rubric weights, model sizes, and training steps—could map the reward landscape that RL policies exploit, informing more robust multi-objective reward design for scientific tasks.

Dynamic Rewards

Static rubric weights can be gamed once the policy identifies the easiest path. Adaptive reward schedules—annealing the composition weight down as stability improves, or automatically reweighting based on which rubric is currently the bottleneck—could keep the gradient signal aligned with actual learning progress throughout training.

DFT-Assisted Initialization

Even if replacing DFT outright remains unrealistic, complementing it is already useful. One practical direction is to use the model as a proposal generator: sample diverse but chemically plausible crystals, then feed those structures in as better initial conformations before running conventional DFT relaxations.

DFT-in-the-Loop Rewards

With more compute, a stronger next step is to run actual DFT relaxations on structures generated by a fine-tuned model and use those outcomes as the reward signal. That would open the door to PPO- or RLHF-like fine-tuning over traces scored by much higher-quality physics, rather than relying only on surrogate rewards from CHGNet and heuristic rubrics.

CIF format: A Crystallographic Information File (CIF) is a standard text format for crystal structures. It typically stores unit-cell lengths and angles, symmetry or space-group metadata, atomic species, fractional coordinates, occupancies, and related crystallographic annotations needed to reconstruct the structure.

📚 Implementation Note: Training hyperparameters (batch size, learning rate, epochs) and rubric weights were adapted from CHGNet's original training protocol and empirically tuned for the RL environment. The formation energy threshold and bond length tolerances follow standard practices in computational materials discovery.

What's Next?

Several open problems are worth pursuing from here:

Vector feedback for iterative learning. Current feedback is scalar ("energy too high"). The training analysis shows the agent can't fix failed structures, only re-roll. Providing per-atom force vectors as textual feedback—e.g., "atom 3 (O) has residual force 1.2 eV/Å along c"—could give the policy the gradient signal it needs for true iterative refinement.
Curriculum over E_f threshold. Start with a wide threshold (0.5 eV/atom) and anneal it toward 0.05 eV/atom as training progresses—analogous to curriculum learning in RL for locomotion tasks. This avoids the sparse reward problem at initialization.
Generalization to novel compositions. The current agent is trained on mp-20 structures. Testing on out-of-distribution materials (e.g., high-entropy alloys, ternary nitrides outside the training set) would probe whether the policy has learned generalizable structural heuristics or is pattern-matching mp-20 CIFs.
Symmetry-aware action space. All atomic coordinates are currently treated independently. Imposing Wyckoff position constraints during generation—enforcing the correct space group symmetry in the CIF—would dramatically reduce the effective action space and should improve sample efficiency.
Hybrid LLM + optimizer pipeline. Use the LLM as a coarse-grained proposer and run a CHGNet-driven BFGS pass on the output for fine-grained relaxation. This combines the generative flexibility of an LLM policy with the convergence guarantees of gradient-based optimization.

⚠️ What Didn't Work

These are the failure modes discovered during development—documented so others don’t repeat them.

1. Equal weighting of composition and stability rewards leads to reward hacking

In an early version of the reward function, composition and formation energy rubrics were given equal weight (~0.40 each). The agent, at different stages of training, would discover that it could achieve consistently higher rewards by memorizing a single structurally simple crystal—a binary oxide with symmetric Wyckoff positions or a bi-metallic crystal—that reliably scores well on stability metrics. This promoted laziness: instead of trying to comply with composition constraints, the model would just output an easy structure. Once this pattern was discovered, the model quickly degenerated from shorter to almost non-existent reasoning trajectories. The fix was to decrease the importance of the composition rubric while keeping it non-zero: low enough that the policy is not rewarded enough to make quick, undesired reward gains, but high enough to provide gradient signal early in training when the model is too weak to pass the stability rubric.

2. Redundant rubrics increase variance without improving performance

The environment originally used both a force rubric and a formation energy rubric to check validity and stability. Both were intended to provide complementary signals about why generation was failing, thereby improving training. In practice, performance never improved when both were active. My hypothesis is that including both rubrics increased reward variance to the point that it was harder for the model to get gradient updates in the right direction—the signals partially cancelled rather than reinforced.

3. Multi-turn feedback can be a re-roll signal, not a repair signal

In multi-turn rollouts, I expected the model to use feedback from a failed first attempt to diagnose and correct structural deficiencies. Across 37 first-turn failures from the 30B run, there were zero physical recoveries—no corrections to formation energy, forces, or bond lengths. Seven cases showed syntactic recoveries (fixing a malformed CIF tag), but the model treated scalar feedback like "Formation energy is too high" as a cue to re-roll rather than to diagnose. It rephrases its chain-of-thought reasoning and outputs a near-identical structure with minor random perturbations. The model can repair a broken file; it cannot repair a bad structure.