Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MDAgent2: Large Language Model for Code Generation and Knowledge Q&A in Molecular Dynamics

Zhuofan Shi, Hubao A, Yufei Shao, Mengyan Dai, Yadong Yu, Pan Xiang, Dongliang Huang, Hongxu An, Chunxiao Xin, Haiyang Shen, Zhenyu Wang, Yunshan Na, Gang Huang, Xiang Jing

TL;DR

MDAgent2 addresses the challenge of producing executable, domain-accurate LAMMPS scripts and MD knowledge QA by building a domain-specific data pipeline and a three-stage post-training (CPT, SFT, RL) to create MD-Instruct and MD-Code. It introduces MD-GRPO, a closed-loop RL framework that uses actual simulation outcomes as rewards to refine code generation, and MDAgent2-RUNTIME, a deployable multi-agent system for automated code generation, execution, evaluation, and self-correction. The paper also presents MD-EvalBench, the first benchmark for MD knowledge QA and LAMMPS code generation, and shows that the MD-series LLMs outperform several baselines on QA tasks and achieve strong executable-code performance with runtime feedback. Together, these contributions demonstrate the feasibility and impact of domain-adapted, feedback-driven LLMs for AI-for-Science and industrial-scale molecular simulations.

Abstract

Molecular dynamics (MD) simulations are essential for understanding atomic-scale behaviors in materials science, yet writing LAMMPS scripts remains highly specialized and time-consuming tasks. Although LLMs show promise in code generation and domain-specific question answering, their performance in MD scenarios is limited by scarce domain data, the high deployment cost of state-of-the-art LLMs, and low code executability. Building upon our prior MDAgent, we present MDAgent2, the first end-to-end framework capable of performing both knowledge Q&A and code generation within the MD domain. We construct a domain-specific data-construction pipeline that yields three high-quality datasets spanning MD knowledge, question answering, and code generation. Based on these datasets, we adopt a three stage post-training strategy--continued pre-training (CPT), supervised fine-tuning (SFT), and reinforcement learning (RL)--to train two domain-adapted models, MD-Instruct and MD-Code. Furthermore, we introduce MD-GRPO, a closed-loop RL method that leverages simulation outcomes as reward signals and recycles low-reward trajectories for continual refinement. We further build MDAgent2-RUNTIME, a deployable multi-agent system that integrates code generation, execution, evaluation, and self-correction. Together with MD-EvalBench proposed in this work, the first benchmark for LAMMPS code generation and question answering, our models and system achieve performance surpassing several strong baselines.This work systematically demonstrates the adaptability and generalization capability of large language models in industrial simulation tasks, laying a methodological foundation for automatic code generation in AI for Science and industrial-scale simulations. URL: https://github.com/FredericVAN/PKU_MDAgent2

MDAgent2: Large Language Model for Code Generation and Knowledge Q&A in Molecular Dynamics

TL;DR

MDAgent2 addresses the challenge of producing executable, domain-accurate LAMMPS scripts and MD knowledge QA by building a domain-specific data pipeline and a three-stage post-training (CPT, SFT, RL) to create MD-Instruct and MD-Code. It introduces MD-GRPO, a closed-loop RL framework that uses actual simulation outcomes as rewards to refine code generation, and MDAgent2-RUNTIME, a deployable multi-agent system for automated code generation, execution, evaluation, and self-correction. The paper also presents MD-EvalBench, the first benchmark for MD knowledge QA and LAMMPS code generation, and shows that the MD-series LLMs outperform several baselines on QA tasks and achieve strong executable-code performance with runtime feedback. Together, these contributions demonstrate the feasibility and impact of domain-adapted, feedback-driven LLMs for AI-for-Science and industrial-scale molecular simulations.

Abstract

Molecular dynamics (MD) simulations are essential for understanding atomic-scale behaviors in materials science, yet writing LAMMPS scripts remains highly specialized and time-consuming tasks. Although LLMs show promise in code generation and domain-specific question answering, their performance in MD scenarios is limited by scarce domain data, the high deployment cost of state-of-the-art LLMs, and low code executability. Building upon our prior MDAgent, we present MDAgent2, the first end-to-end framework capable of performing both knowledge Q&A and code generation within the MD domain. We construct a domain-specific data-construction pipeline that yields three high-quality datasets spanning MD knowledge, question answering, and code generation. Based on these datasets, we adopt a three stage post-training strategy--continued pre-training (CPT), supervised fine-tuning (SFT), and reinforcement learning (RL)--to train two domain-adapted models, MD-Instruct and MD-Code. Furthermore, we introduce MD-GRPO, a closed-loop RL method that leverages simulation outcomes as reward signals and recycles low-reward trajectories for continual refinement. We further build MDAgent2-RUNTIME, a deployable multi-agent system that integrates code generation, execution, evaluation, and self-correction. Together with MD-EvalBench proposed in this work, the first benchmark for LAMMPS code generation and question answering, our models and system achieve performance surpassing several strong baselines.This work systematically demonstrates the adaptability and generalization capability of large language models in industrial simulation tasks, laying a methodological foundation for automatic code generation in AI for Science and industrial-scale simulations. URL: https://github.com/FredericVAN/PKU_MDAgent2
Paper Structure (25 sections, 4 equations, 4 figures, 4 tables)

This paper contains 25 sections, 4 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Results
  3. Evaluation Benchmark for Molecular Dynamics: MD-EvalBench
  4. Experimental Setup
  5. Evaluation Protocol
  6. Baselines and Compared Models
  7. Evaluation of QA Ability
  8. Evaluation of Code Generation Ability
  9. Methods
  10. Construction of Training Datasets
  11. MD-Knowledge: A Pretraining Corpus for Molecular Dynamics
  12. MD-InstructQA: A Domain-Specific Instruction–Answer Dataset for Molecular Dynamics
  13. MD-CodeGen: A Code Generation Dataset for Molecular Dynamics Simulations
  14. Post-Training for MD-LLM
  15. CPT and SFT
  16. ...and 10 more sections

Figures (4)

  • Figure 1: Overall workflow of the proposed MDAgent2 framework, integrating data construction, model training, multi-agent runtime, and evaluation.
  • Figure 2: Comparison of Code-Score-Human and Execution Success@3 across all methods.
  • Figure 3: Overview of the MD-GRPO training framework. The policy LLM generates candidate codes through the rollout module, which are evaluated by a reward model against a reference LLM. Group computation aggregates rewards and updates the policy via KL-regularized optimization. Low-scoring tasks are rewritten for robustness.
  • Figure 4: Overall architecture of the MDAgent2-RUNTIME multi-agent system.