Skip to content

MACS: Multi-Agent Reinforcement Learning for Optimization of Crystal Structures

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

MACS frames periodic crystal structure geometry optimization as a partially observable Markov game where each atom is an independent agent, trained via independent Soft Actor-Critic (SAC) to collectively minimize atomic forces. On six diverse chemical compositions, MACS achieves 34% lower wall-clock time and 28% fewer energy calculations than the strongest classical baseline (BFGSLS), with the lowest failure rate (0.36%), and generalizes zero-shot to unseen compositions and larger structures.

Problem

Classical geometry optimizers (BFGS, FIRE, CG) for periodic crystal structures require many energy evaluations—each expensive—or many steps to converge, creating a bottleneck in crystal structure prediction (CSP) workflows that require thousands of local optimization runs. Prior learning-to-optimize approaches addressed finite molecular clusters but not periodic crystals, which introduce unit-cell periodicity, boundary crossing, and variable atom counts.

Method

MACS models crystal geometry optimization as a partially observable Markov game (POMG): each atom in the unit cell is an agent with a local observation vector (Eq. 1–2) containing its covalent radius, a gradient-norm-derived action scaling factor, current CHGNet gradient vector, log gradient norm, gradient change from previous step, and previous displacement, plus the feature vectors and relative positions of its 12 nearest neighbors (obs. length 204). Actions are displacements scaled by min(|grad|, c_max) (Eq. 3–4), enabling smooth periodic boundary crossing without positional coordinates. The individual reward is log|g_t| − log|g_{t+1}| (Eq. 5), a logarithmic ratio of gradient norms that avoids over-weighting early large-force reductions. Training uses independent SAC with shared policy and Q-networks (two-layer MLP + ReLU) across all atom types, a replay buffer of 10M transitions, and discount γ = 0.995. CHGNet serves as the energy/force oracle throughout.

Key Contributions

  • First application of MARL to periodic crystal structure optimization.
  • Novel POMG formulation treating atoms as decentralized agents with local observations, enabling natural scalability without explicit atom-count encoding.
  • Gradient-norm-scaled action space and log-ratio reward function, both shown via ablation to be critical to sample efficiency and convergence.
  • Zero-shot transferability to unseen compositions (Sr-Ti-O family) and structures up to 2× training size without retraining.

Results

  • Average across all test sets: MACS Tmean = 66s vs. BFGSLS (best baseline) 152s — 34% faster; Cmean = 171 vs. 253 energy calculations — 28% fewer; failure rate PF = 0.36% vs. BFGSLS 0.36%, CG 18.22%, MDMin 46.19%.
  • Steps (Nmean): BFGSLS requires ~5% fewer steps than MACS on average, but BFGSLS uses multiple energy calls per step, so total energy calculations and time are higher.
  • Scalability: Performance advantage holds as atom count scales from K → 2K training-size atoms across all compositions.
  • Zero-shot transfer (Sr₂TiO₄, Sr₃Ti₂O₇, Sr₄Ti₃O₁₀): MACS outperforms all baselines in Tmean and Cmean on both same-size and 2× structures.
  • Energy quality: Optimized structure energy distributions match those of classical baselines — MACS finds equivalent local minima, not shallower ones.
  • Ablation: Removing force-related features (feat.6) causes complete failure; removing history features degrades Nmean by ~50%; gradient-norm action scaling is critical for sample efficiency.

Limitations

  • Atomic species distinguished only by covalent radius; richer descriptors (e.g., electronegativity, oxidation state, existing ML descriptors) are unexplored.
  • Unit cell vectors (lattice parameters) are fixed during optimization; full variable-cell relaxation is not yet supported.
  • Trained and evaluated exclusively with CHGNet as the energy oracle; accuracy is bounded by CHGNet's DFT-training distribution.
  • Policy architecture uses MLPs without recurrence; ablations suggest history features matter, implying RNNs could improve performance further.
  • Training compositions are limited to six; generalization to arbitrary chemistry across the full periodic table is unvalidated.

Relevance to Agentic AI / LLM Agents

MACS is a clean demonstration of the decentralized multi-agent coordination paradigm applied to a scientific optimization task: each atom-agent acts on local observations to produce emergent global behavior, directly instantiating the cooperative MARL setting central to agentic AI research. The zero-shot transfer result — a single trained policy generalizing to novel compositions and larger state spaces without retraining — directly addresses a key open problem in agent generalization. The independent SAC with shared networks pattern parallels architectural choices in scalable LLM-based agent systems (shared backbone, decentralized execution). While MACS agents are not LLM-based, its problem formulation (partial observability, emergent cooperation, scalability to variable numbers of agents) is instructive for designing agentic systems where the number of autonomous actors is not fixed at training time.