CrystalReasoner: Reasoning and RL for Property-Conditioned Crystal Structure Generation

CrystalReasoner (CrysReas) is an end-to-end LLM framework for generating crystal structures from natural language instructions. It uses supervised fine-tuning (SFT) to teach crystal-structure generation, thinking traces to introduce crystallographic and physical priors before coordinates, and reinforcement learning (RL) with verifiable rewards to improve validity, stability, and property conditioning. Please see our work at crystalreasoner.github.io.

The first model is our main model. The subsequent models were optimized for property conditioning. You can directly use them to inference without installing our repo. Our repo is for training and evaluation. See how to use them to inference in the huggingface pages.

The easiest way to install prerequisites is via uv, a fast Python package and project manager.

Assumes you are running Linux and have a CUDA GPU. The CrysReas environment can be installed via the following command. Notably, the torch version will be 2.8.0+cu128 because of Verl, and we adjust other modules to fit its torch version.

# Better for building
export MAX_JOBS=8

# Initialize the environment
uv venv .venv --python 3.12
source .venv/bin/activate

# Basics
uv pip install pip
uv pip install -e .
uv pip install -r requirements.txt
uv pip install ninja

# Install Verl
bash submodules/verl/scripts/install_vllm_sglang_mcore.sh

# Install torch-sim
uv pip install -e "submodules/torch-sim[mace,test]"

# Install mattergen
uv pip install torch-cluster -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
uv pip install torch-scatter -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
uv pip install torch-sparse -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
uv pip install -e submodules/mattergen

# Some fix, you may not need them? But if you encounter problems, try to use them.
uv pip install e3nn==0.4.4
uv pip install mattersim==1.2.0
uv pip install "numpy<2.2"
uv pip install "ase==3.25.0"
uv pip install -U mp-api emmet-core
uv pip install pymatgen==2024.10.29

Optionally, you can install apex for better training and inference.

# Clone apex somewhere and install.
export FORCE_CUDA=1
git clone https://github.com/NVIDIA/apex.git
cd apex
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .

Please download assets/MP/MP_shelve.parquet via Github LFS (git lfs pull) or HF Hub, and convert it to python shelve dataset by python scripts/parquet_to_shelve.py, which contains all the data we use in training and inference as ground truth.

See doc/DATASET.md to understand how this dataset is created from Materials Project. See doc/METRICS.md for the metric names and required inputs.

MatterGen's MP2020 correction file is also required for energy-above-hull evaluation:

curl -L -o submodules/mattergen/data-release/alex-mp/reference_MP2020correction.gz \
  https://github.com/microsoft/mattergen/raw/main/data-release/alex-mp/reference_MP2020correction.gz

You can follow these instructions to recreate our checkpoints and run our experiments.

scripts/run.py is the main entrypoint. It reads experiment definitions from scripts/config.py.

python scripts/run.py [-h] [--experiment STR] [--output STR] [--dry_run | --no_dry_run] [extra overrides...]

It supports SFT, RL/GRPO, generation, checkpoint merge, and evaluating metrics through one registry.

Before training, you must set environment variables SWANLAB_API_KEY and HF_TOKEN. Notably, you can write .env file in the project folder to do this.

In default the output of this script is located at logs/{experiment_name}.out and will not be printed directly to the command line. To output to the terminal, set --output=terminal. To use a timestamped log file, set --output=time.

Common examples:

# Supervised fine-tuning baselines
python scripts/run.py sft_no_thinking
python scripts/run.py --output=terminal sft_no_thinking
python scripts/run.py sft_thinking
python scripts/run.py sft_thinking gpu=2

# Merge FSDP checkpoints into Hugging Face format
python scripts/run.py merge \
  --path checkpoints/thinking/global_step_1514 \
  --output-path checkpoints_merged/thinking

# Generate samples from a merged checkpoint
python scripts/run.py generate checkpoints_merged/thinking

# RL/GRPO alignment from a merged SFT checkpoint
python scripts/run.py rl_thinking_novel

# Run metric_process through the unified runner
python scripts/run.py run_metric \
  --path checkpoints_merged/thinking/conditional+thinking.parquet \
  --metrics-name composition_consistency spacegroup_consistency stable_unique_novel \
  --prompt-type conditional+thinking \
  --forced

Representative registered jobs include:

• SFT: sft_no_thinking, sft_thinking, sft_crystaltextllm, sft_crystaltextllm_8, sft_plaid_wyckoff, sft_plaid_wyckoff_8
• RL/GRPO: rl_no_thinking, rl_thinking, rl_thinking_novel, rl_spacegroup_thinking, rl_elastic_thinking_new, rl_cte_thinking
• Generation: generate, generate_rl, generate_elastic, generate_cte, generate_crystaltextllm, generate_plaid_wyckoff
• Metrics: run_metric, metric_qha, metric_elastic
• Checkpoint conversion: merge

See doc/TRAINING.md for training details.

The preferred evaluation path is crysreas.metric_process. python scripts/run.py run_metric is a copy of this command. It is a DataFrame-based pipeline with named metrics, cache-like ensure_* helpers, optional multiprocessing for CPU metrics, and serialization helpers for Parquet files containing structures.

Example:

python -m crysreas.metric_process \
  --path checkpoints_merged/thinking/conditional+thinking.parquet \
  --metrics-name simple_structure structure_validity smact_validity composition_consistency spacegroup_consistency stable_unique_novel \
  --prompt-type conditional+thinking \
  --num-workers 1 \
  --forced

Important metric groups include:

• Parsing and validity: simple_structure, structure_validity, smact_validity
• Ground-truth consistency: gt, composition_consistency, spacegroup_consistency
• Relaxation and stability: relaxed_structures, energy_above_hull, stable_unique_novel
• Property rewards: energy_reward, smooth_energy_reward, elastic_reward, elastic_reward_all
• Property calculations: fmax, elastic_properties, cte

Use doc/METRICS.md as the source of truth for available metrics, column dependencies, prompt-type configuration, multiprocessing behavior, and Parquet structure serialization.

After training, the checkpoint is genreated in checkpoints/. After merging, it is in checkpoints_merged/.

Precomputed experiment checkpoints are available at huggingface, will the sample outputs are available in the checkpoints_merged folder of this repo. Below is the relation between the parquets in checkpoints_merged folder and the huggingface models:

See more in doc/CHECKPOINTS.md.

The default configuration assumes:

• At least 1 GPU for SFT.
• At least 2 GPUs for RL/GRPO jobs.
• At least 1 GPU for MLIP-backed reward or metric calculations.
• At least 4 CPUs cores for typical training and evaluation jobs.

Some registered jobs request more CPU memory and CPU cores in scripts/config.py, especially elastic and CTE tasks.

If you find our work useful, please cite:

@article{wu2026crysreas,
  title={CrystalReasoner: Reasoning and RL for Property-Conditioned Crystal Structure Generation},
  author={Yuyang Wu and Stefano Falletta and Delia McGrath and Sherry Yang},
  year={2026},
  journal={arXiv preprint arXiv:2605.14344},
  url={https://arxiv.org/abs/2605.14344}
}