MDAgent2: A Large Model Framework for Code Generation and Knowledge Q&A in the Molecular Dynamics Domain

The Peking University team released MDAgent2, the first end-to-end large language model framework in the molecular dynamics domain, supporting automatic LAMMPS script generation and domain knowledge Q&A. Through three-stage post-training (continuous pre-training, supervised fine-tuning, reinforcement learning) and the innovative MD-GRPO closed-loop reinforcement learning mechanism, it significantly improves code executability, lowers the threshold for MD simulations, and promotes the application of AI for Science in industrial-level simulation scenarios.

Molecular Dynamics (MD) simulation is a core research method in materials science, biophysics, and chemical engineering. However, writing high-quality LAMMPS scripts requires deep domain knowledge and practical experience, and this threshold limits the popularization of the technology. General-purpose large language models face three major challenges in this domain: scarcity of high-quality domain data, high deployment costs for closed-source models, and lack of executability verification for generated code. MDAgent2 is designed to address these pain points.

MDAgent2 adopts a modular multi-agent architecture, including two domain models: MD-Instruct (knowledge Q&A and instruction understanding) and MD-Code (LAMMPS code generation). Three high-quality datasets are constructed: MD-Knowledge (basic theories and concepts), MD-InstructQA (structured Q&A pairs), and MD-CodeGen (natural language-LAMMPS script pairs), providing a solid foundation for model training.

A progressive three-stage training approach is adopted: 1. Continuous Pre-Training (CPT): Unsupervised learning on MD-related literature and code to build domain language understanding; 2. Supervised Fine-Tuning (SFT): Training with annotated instruction data and code pairs to enable natural language-to-LAMMPS command conversion and answering professional questions; 3. Reinforcement Learning (RL): The innovative MD-GRPO method uses simulation results as reward signals to form closed-loop feedback.

MD-GRPO integrates the LAMMPS simulator into the training loop, fusing code generation and execution verification, focusing on the actual execution behavior of the code. A key feature is low-reward trajectory recycling: analyzing error patterns of failed scripts and extracting intermediate information for subsequent training to improve sample efficiency.

The MDAgent2-RUNTIME deployment system is developed, integrating modules for code generation, automatic execution, result evaluation, and self-correction. Users can complete the entire process by describing their needs in natural language. The MD-EvalBench benchmark test set is constructed, including code generation and Q&A tasks, to evaluate code executability and physical rationality. Tests show that MDAgent2 outperforms strong baseline models.

It provides usage interfaces such as Python API and Web applications, supports Docker deployment, and open-sources training processes and fine-tuning scripts. In the future, it is expected to expand to more scientific computing fields such as quantum chemistry, helping researchers focus on the design and analysis of physical problems and promoting educational popularization. The open-source release promotes domain development, and more participation is expected.