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

Zhuofan Shi; Hubao A; Yufei Shao; Dongliang Huang; Hongxu An; Chunxiao Xin; Haiyang Shen; Zhenyu Wang; Yunshan Na; Gang Huang; Xiang Jing

Source-linked AI summary

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

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

arXiv:2601.02075v4Jan 5, 2026cs.CEcs.LG

TL;DR

LAMMPS-based MD simulations require specialized, time-consuming scripting, and adapting LLMs to these tasks remains challenging. MDAgent2 combines domain-specific datasets and models with execution-feedback training and an automated runtime. Experiments report capabilities across MD QA and code generation, and the authors identify limited task coverage as a boundary of the current dataset.

Problem
LAMMPS simulation scripts require domain expertise and extensive manual work, while adapting LLMs to specialized MD tasks remains challenging.
Source: §1 · p. 2
Method
MDAgent2 combines MD datasets and domain-adapted models with MD-GRPO execution-reward training and a runtime integrating code generation, execution, evaluation, and self-correction.
Source: §1 · p. 4, §3.2.2 · p. 10, §4 · p. 14
Results
Experiments on MD-EvalBench demonstrate that MDAgent2 and the MD-series LLMs address MD question answering and simulation-code generation.
Source: §1 · p. 5, §4 · p. 14
Takeaways & Limitations
MDAgent2 advances MD code generation beyond one-shot text-to-code toward a closed-loop workflow.
Source: §4 · p. 14
The dataset currently covers thermodynamic-property, fluid-dynamics, and materials mechanical-property simulations, with broader task coverage planned.
Source: Limitation and Outlook · p. 15

Abstract

from arXiv · show

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

1 Introduction

LAMMPS-based molecular dynamics scripting demands domain expertise and extensive manual work, while MD-specific data, benchmarks, and closed-loop code optimization remain limited. MDAgent2 addresses these needs with domain datasets and models, execution-feedback training, and an automated multi-agent runtime.

1 Introduction

LAMMPS scripts require domain expertise and hands-on experience, and repetitive manual workflows can bottleneck research efficiency.The scripts also involve complex models, physical constraints, and structured input formats.Source: §1 · p. 2
The paper identifies gaps in MD-specific high-quality data and benchmarks, closed-loop code optimization, and locally deployable models.Existing benchmarks have limited MD and LAMMPS coverage, while proprietary and ultra-large models constrain practical deployment.Source: §1 · pp. 2–3
MDAgent2 combines MD knowledge Q&A and simulation-code generation in an end-to-end framework.The framework is designed for both tasks in molecular dynamics.Source: §1 · p. 4
Three datasets span MD knowledge, instruction QA, and LAMMPS code generation across material systems, physical conditions, and simulation tasks.They are MD-Knowledge, MD-InstructQA, and MD-CodeGen.Source: §1 · p. 4
CPT, SFT, and RL form the post-training pipeline for adapting the models to MD knowledge and task requirements.CPT injects MD and LAMMPS knowledge, while SFT aligns models with MD tasks.Source: §1 · p. 4
MD-GRPO uses simulation execution and evaluation scores as rewards, recycling low-reward trajectories to refine code generation.The method is intended to improve code executability and physical correctness.Source: §1 · p. 4
MD-Instruct targets domain understanding, while MD-Code targets simulation-code generation.Both models are based on the Qwen3 series.Source: §1 · p. 5
MDAgent2-RUNTIME automates code generation, execution, evaluation, and self-correction from natural-language task descriptions.The paper reports strong QA and code-generation capabilities against selected baselines, and improved correct, executable LAMMPS code with the runtime.Source: §1 · p. 5

2 Results

MD-EvalBench assesses MD knowledge, LAMMPS syntax, and code generation, with results reported for QA and code-generation benchmarks. The reported findings include improved execution success with MDAgent2-RUNTIME and competitive QA performance from MD-Instruct-8B.

2.1 Evaluation Benchmark for Molecular Dynamics: MD-EvalBench

MD-EvalBench comprises MD-KnowledgeEval, LAMMPS-SyntaxEval, and LAMMPS-CodeGenEval for theoretical knowledge, syntax comprehension, and code generation.MD-KnowledgeEval has 336 questions and LAMMPS-SyntaxEval has 333; code-generation samples are evaluated structurally and functionally.Source: §2.1 · p. 5

Evaluation Protocol

Experiments are repeated three times and averaged; Exec-Success@k measures successful execution among k candidates, while Code Human Score rates readability, robustness, and physical correctness on a 0–10 scale.The compared systems include MDAgent, Direct Prompting, and Qwen3.Source: Evaluation Protocol · p. 6, Baselines and Compared Models · p. 6

2.3 Evaluation of QA Ability

MD-Instruct-8B scores 74.67 on average, surpassing Qwen-Flash and Qwen3-14B and narrowing the gap to Qwen3-32B.Qwen3-Max has the highest average score, 82.49.Source: §2.3 · pp. 6–7, Table 1 · p. 6

2.4 Evaluation of Code Generation Ability

ExecSucc@3 rises from 14.23 % to 37.95 % for MD-Code-8B with MDAgent2-RUNTIME, while Code-Score-Human changes from 9.29 to 9.32.The reported gains are associated with iterative evaluation, self-correction, and LAMMPS-specific tools that provide actionable feedback.Source: Figure 2 · p. 7, §2.4 · pp. 7–8
Across three evaluated backbone models, MDAgent2 consistently outperforms MDAgent in code generation.The paper attributes the improvement to LAMMPS-tailored tools and actual execution feedback unavailable in MDAgent.Source: §2.4 · p. 8
After domain-specific post-training, MD-Code-8B reaches Code-Score-Human 9.29 under Direct Prompting, improving substantially over the Qwen3-8B baseline.The paper attributes the gain to the full post-training stack, including execution-feedback optimization in MD-GRPO.Source: §2.4 · p. 8

3.1 Construction of Training Datasets

The training datasets are constructed for incremental pre-training, fine-tuning, and reinforcement learning. Their processes combine collected domain sources, automated transformation and generation, and expert review.

3.1 Construction of Training Datasets

Three datasets were constructed for incremental pre-training, fine-tuning, and reinforcement learning.They support multi-task capability training and evaluation.Source: §3.1 · p. 8

MD-Knowledge: A Pretraining Corpus for Molecular Dynamics

MD-Knowledge draws on thousands of MD-related papers, textbooks, technical documents, and public manuals, then applies multi-stage cleaning.The pipeline includes filtering, deduplication, and quality evaluation.Source: MD-Knowledge: A Pretraining Corpus for… · pp. 8–9

MD-InstructQA: A Domain-Specific Question–Answer Dataset for Molecular Dynamics

MD-InstructQA is generated from source documents through structured PDF-to-Markdown conversion, semantic chunking, domain-label question generation, and source-aligned answers.Questions and answers are paired in Alpaca format, with selectable chain-of-thought in answer generation.Source: MD-InstructQA: A Domain-Specific… · p. 9

MD-CodeGen: A Code Generation Dataset for Molecular Dynamics Simulations

MD-CodeGen pairs task descriptions with LAMMPS code, using expert-collected examples and template-generated tasks that are manually checked.Simulation experts review and refine scripts generated with MDAgent2-RUNTIME and SOTA LLMs.Source: MD-CodeGen: A Code Generation Dataset… · pp. 9–10

3.2 Post-Training for MD-LLM

MDAgent2 post-trains MD models through CPT, SFT, and reinforcement learning, culminating in MD-GRPO’s simulation-feedback loop for refining LAMMPS code. Its reward design combines output-format compliance with multidimensional script-quality assessment.

3.2.1 CPT and SFT

CPT mixes MD-Knowledge with general-domain data, then SFT mixes MD-InstructQA with general instructions and seeds code mappings with MD-CodeGen samples.These stages produce MD-Instruct-8B, described as capable of answering MD questions with strong domain fidelity.Source: §3.2.1 · p. 10

3.2.2 MD-GRPO

MD-GRPO executes candidate LAMMPS code in simulation, scores code and results, and uses low-scoring trajectories to refine later training queries.This creates a trajectory-aware, reward-driven self-improvement loop for learning valid code and correcting observed failures.Source: §3.2.2 · pp. 10–11

Total Reward

The description defines format reward Rformat as 1 only when required formatting constraints are met (otherwise 0), correctness reward R̂correct as a normalized [0,1] LAMMPS-script quality score, and defaults λ1=1, λ2=5.Source: Total Reward · p. 11

Format Reward

The binary format reward checks ordered <think> and <answer> blocks and requires the answer to be valid JSON with exactly the task-required fields.Training checks format with regular expressions and strict JSON schema validation.Source: Format Reward · pp. 11–12

Correct Reward

The correctness reward combines additive bonuses and penalties across dimensions of LAMMPS script quality, then maps the raw score to [0, 1].The dimensions are based on MDAgent’s evaluation framework and were refined and expanded with domain experts.Source: Correct Reward · p. 12

3.3 Runtime Multi-Agent Framework

MDAgent2-RUNTIME automates a workflow from natural-language simulation tasks through LAMMPS code generation, execution, and evaluation. Its agents iteratively use tool checks and evaluation feedback to refine code, with optional human input.

3.3 Runtime Multi-Agent Framework

MDAgent2-RUNTIME automatically generates LAMMPS code, executes simulations, and evaluates results from a natural-language task description.The system is organized as Code Generator, Code Runner, and Result Evaluator.Source: §3.3 · p. 13, Figure 4 · p. 13
The Code Generator iteratively revises code using syntax checks and potential-file tools, stopping at convergence or a predefined iteration limit.Potential-file tools can recommend similar files when a specified potential is unavailable.Source: §3.3 · pp. 13–14
The Code Runner executes LAMMPS code in a sandboxed Docker container and stores execution results with an automatically generated summary.The sandbox is described as supporting environment isolation and reproducibility.Source: §3.3 · p. 14
The Result Evaluator scores code and simulation outputs, returns reward and penalty feedback, and triggers regeneration when the score is below a threshold.Its criteria assess multiple dimensions of code and simulation quality.Source: §3.3 · p. 14
Users can pause the workflow to modify code or parameters, or provide domain-specific instructions before execution.The runtime also provides visualization, quality evaluation, syntax verification, and potential-file tools.Source: §3.3 · pp. 13–14

4 Discussion

MDAgent2 combines domain-specific data, adapted models, simulation-feedback reinforcement learning, and a deployable runtime for MD question answering and code generation. Experiments support its effectiveness on MD-EvalBench, while the authors identify task coverage and multimodal evaluation as areas for future work.

4 Discussion

The framework combines a domain-specific data pipeline, MD-Instruct and MD-Code, MD-GRPO using simulation outcomes as rewards, and a runtime integrating generation, execution, evaluation, and self-correction.The authors describe the workflow as moving MD code generation beyond one-shot text-to-code.Source: §4 · p. 14
Experiments on MD-EvalBench demonstrate that MDAgent2 and its MD-series models address molecular dynamics question answering and simulation code generation.The paper also describes future work to expand supported task types and incorporate visual simulation outputs as evaluation inputs.Source: §4 · p. 14
The authors propose that the data-construction paradigm, execution-aware reinforcement learning, and closed-loop workflow may apply to other simulation-driven AI-for-Science scenarios.Molecular dynamics is the representative domain studied in this work.Source: §4 · p. 15

Limitation and Outlook

Current dataset coverage includes thermodynamic-property, fluid-dynamics, and materials mechanical-property simulations, with expansion to additional task types planned.Source: Limitation and Outlook · p. 15
Future work will explore multimodal LLM integration, using visualized thermodynamic curves, atomic trajectories, and structural evolution as additional evaluation inputs.The authors propose these inputs to improve interpretability and assessment granularity.Source: Limitation and Outlook · p. 15
The authors describe the framework as general and scalable, with modular design intended to facilitate adaptation across molecular dynamics and other industrial and materials-science simulation domains.Source: Limitation and Outlook · p. 15

Appendix A Tools Details

Appendix A presents Table A1 as an inventory of tools integrated into MDAgent2 and their functions in the molecular dynamics workflow.

Appendix A Tools Details

Table A1 lists tools integrated into MDAgent2 and their specific functions in the molecular dynamics workflow.Source: Appendix A · p. 18, Table A1 · p. 19
The appendix positions the tool inventory as a supplement to the system architecture described in Section 3.Source: Appendix A · p. 18
The table's stated scope is tool functionality; the supplied passage does not specify individual tool names or operations.Source: Table A1 · p. 19

Appendix B Evaluation Details

Appendix B describes a multi-dimensional rubric for assessing generated LAMMPS scripts, used both for expert ground-truth evaluations and automated scoring prompts.

Appendix B Evaluation Details

The evaluation framework covers eight correctness dimensions, ranging from syntax correctness to physical soundness.Source: Appendix B · p. 19, Table B2 · p. 20
The rubric provides human experts with a unified guideline for ground-truth evaluation and supplies dimensions and deduction rules for the LLM scoring prompt.Source: Appendix B · p. 19
The authors state that aligning the LLM scoring instructions with the expert-verified rubric keeps the automated correctness reward consistent with human scientific judgment.Source: Appendix B · p. 19

Appendix C Datasets Details

Appendix C reports statistical analysis of the study's datasets and identifies tables covering dataset statistics and evaluation-set question types and difficulty levels.

Appendix C Datasets Details

The authors report comprehensive statistical analysis of all datasets referenced in the work.Source: Appendix C · p. 20
Table C3 presents statistics for the training and evaluation datasets.Source: Table C3 · p. 21
Table C4 covers the distribution of question types and difficulty levels in the evaluation datasets.Source: Table C4 · p. 21

Appendix D Case Study

The case study follows MDAgent2-RUNTIME generating and revising LAMMPS code for simulating the melting of a Cu–Ni nanoparticle. The trajectory includes potential-file correction, code checks, execution, and result-based requests for further revision.

The case study targets a LAMMPS simulation of melting in a Cu–Ni nanoparticle.Source: Appendix D · p. 21
When the script references unavailable CuNi.eam, the potential checker recommends the compatible CuNi.eam.alloy file, which the revised script uses.Source: Appendix D · p. 22
The revised script passes potential-availability and LAMMPS syntax checks before being run.Source: Appendix D · p. 22
The evaluator returns a final score of 4, below the acceptance threshold, and marks the result unacceptable, requiring further code revision.Source: Appendix D · p. 22
The revised code defines a temperature loop from 900 in increments of 10 and records temperature and per-atom energy during NVT runs.Source: Appendix D · p. 24