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

Elena Zamaraeva, Christopher Collins, George R Darling, Matthew Stephen Dyer, Bei Peng, Rahul Savani, Dmytro Antypov, Vladimir Gusev, Judith Clymo, Paul G. Spirakis, Matthew Rosseinsky

Keywords: multi-agent reinforcement learning, learning to optimize, geometry optimization, local optimization, atomic structure optimization, crystal structure

TL;DR: MACS: a new multi-agent reinforcement learning method for faster geometry optimization of crystal structures.

Abstract: Geometry optimization of atomic structures is a common and crucial task in computational chemistry and materials design. Following the learning to optimize paradigm, we propose a new multi-agent reinforcement learning method called Multi-Agent Crystal Structure optimization (MACS) to address the problem of periodic crystal structure optimization. MACS treats geometry optimization as a partially observable Markov game in which atoms are agents that adjust their positions to collectively discover a stable configuration. We train MACS across various compositions of reported crystalline materials to obtain a policy that successfully optimizes structures from the training compositions as well as structures of larger sizes and unseen compositions, confirming its excellent scalability and zero-shot transferability. We benchmark our approach against a broad range of state-of-the-art optimization methods and demonstrate that MACS optimizes periodic crystal structures significantly faster, with fewer energy calculations, and the lowest failure rate.

Primary Area: Reinforcement learning (e.g., decision and control, planning, hierarchical RL, robotics)

Submission Number: 19525