cs.AIOct 24, 2026

LegoSR: A Unified Framework for Self-improving Agent-Guided Scientific Discovery

Authors: Hao Mao, Xu T. Liu, Peng Zhao, Shuai Lu, Yuntian Chen

Organizations: The Hong Kong Polytechnic University, Hong Kong SAR, China · Eastern Institute of Technology, Ningbo, China · Amazon Web Services, Seattle, USA · University of Washington, Seattle, USA · 

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Equation discovery solvers often use different data interfaces, equation representations and evaluation criteria, making their results difficult to compare. We present LegoSR, a unified framework for algebraic symbolic regression and differential equation discovery. LegoSR integrates seven solvers through a common equation representation and evaluation process, and supports searches constrained by partially known equations. Researchers can run solvers directly or use LegoSR Agent, which selects solvers and adapts searches based on solver knowledge, intermediate results and lessons from earlier tasks. On noisy ordinary and partial differential equation benchmarks, LegoSR Agent recovers nearly as many equations as the union of seven expert-tuned solvers and outperforms a generic agent with the same model, tools and budget by up to 29 percentage points. Three examples demonstrate basic use, adaptive solver selection and discovery guided by partial physical knowledge.

Figures & tables

MethodApproachKey characteristics
SINDy/PySINDy [ 2 , 3 , 15 ]Sparse regressionSparse selection from a user-defined library of candidate terms
DLGA-PDE [ 4 ]Neural network and genetic algorithmNeural-network derivative estimation and evolutionary construction of candidate terms
SGA-PDE [ 5 ]Genetic programmingOpen-form discovery of compound and fractional terms using expression trees
PySR [ 6 ]Evolutionary symbolic regressionJoint search over expression structures and numerical constants, balancing accuracy and complexity
DISCOVER [ 7 ]Reinforcement learningSymbol-by-symbol equation generation with constraints on the search space
LLM4ED [ 9 ]Large language modelGeneration and refinement of candidate equations using prompts and evaluation feedback
Table 1: Equation discovery methods integrated in LegoSR and the main characteristics of their original designs.
Figure 1: Architecture of LegoSR. Shared scientific computing modules support heterogeneous solvers, with Equation IR, Sketches and experiment management connecting their workflows.
Figure 2: Architecture of LegoSR Agent. The LLM controller reads knowledge and memory, chooses the next experiment on the LegoSR platform and reads back the common evaluation. A researcher can add prior knowledge as a Sketch with the task.
Figure 3: How LegoSR Agent improves. (a) Within a task, each round rejects terms ( ×\times ), anchors supported terms (filled) and leaves the rest as a hole R\mathcal{R} until the hole is filled (blue) and the equation is checked. (b) Across tasks, each graded task adds a lesson on what worked ( ✓\checkmark ) or failed ( ×\times ), and later tasks read the library.
NoiseUNIONGeneric AgentLegoSR Agent
0%52.2% (70)50.0% (67)48.5% (65)
1%49.3% (66)48.5% (65)50.0% (67)
5%44.8% (60)44.0% (59)45.5% (61)
Table 2: Recovery coverage on 134 algebraic SR problems. Numbers of recovered problems are shown in parentheses.
TaskNoiseUNIONGeneric AgentLegoSR Agent
ODE (40)0%70.0% (28)65.0% (26)67.5% (27)
1%55.0% (22)47.5% (19)52.5% (21)
5%42.5% (17)30.0% (12)40.0% (16)
PDE (17)0%100.0% (17)94.1% (16)94.1% (16)
1%94.1% (16)58.8% (10)88.2% (15)
5%82.4% (14)52.9% 0 (9)82.4% (14)
Table 3: Recovery coverage on 40 ODE and 17 PDE problems. Numbers of recovered problems are shown in parentheses.
Figure 4: Burgers equation recovered with LegoSR. (a) Observed field. (b) Solution of the recovered equation, integrated from t=0t=0 , minus the data.
Figure 5: KdV equation with 1% noise. (a) uxxxu_{xxx} coefficient after each search of LegoSR Agent (blue) and the generic agent (grey); crosses mark wrong equation structures (red: LegoSR Agent), and the band marks ±5%\pm 5\% . (b) Equation at each stage of LegoSR Agent and the best search of the generic agent (grey); numbers are the searches in (a). (c) Data and the final equation integrated from t=0t=0 . (d) Profiles at t=1t=1 .
Figure 6: Chafee–Infante equation. (a) uxxu_{xx} at t=0.26t=0.26 from the surrogate left by the failed stage-3 runs, with the noise-free field for reference. (b) Discovery stages for the first of three runs; crosses mark wrong structures. (c) Normalized residual JJ after dropping each term of the stage-4 equation, relative to JJ of the full equation. (d) Coefficients, normalized by the true values, after the Sketch search (stage 5) and the weak-form fit (stage 6); the band marks ±5%\pm 5\% . (e, f) Data and the recovered equation integrated from the noise-free initial state.
Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

SystemEquations
Harmonic oscillatorx˙=v\dot{x}=v , v˙=−1.69 x\dot{v}=-1.69\,x
Damped harmonic oscillatorx˙=v\dot{x}=v , v˙=−1.69 x−0.26 v\dot{v}=-1.69\,x-0.26\,v
Damped pendulumθ˙=v\dot{\theta}=v , v˙=−1.44sin⁡θ−0.08 v\dot{v}=-1.44\sin\theta-0.08\,v
Duffing oscillatorx˙=v\dot{x}=v , v˙=−x−0.8 x3−0.12 v\dot{v}=-x-0.8\,x^{3}-0.12\,v
Van der Pol oscillatorx˙=v\dot{x}=v , v˙=2 (1−x2) v−x\dot{v}=2\,(1-x^{2})\,v-x
Lorenz systemx˙=10 (y−x)\dot{x}=10\,(y-x) , y˙=x (28−z)−y\dot{y}=x\,(28-z)-y , z˙=xy−83z\dot{z}=xy-\tfrac{8}{3}z
Appendix
Table B1: ODE systems.
ProblemEquationDomainGrid
Burgers (sine initial state) [ 5 ]ut=−uux+0.1 uxxu_{t}=-uu_{x}+0.1\,u_{xx}[−8,8)×[0,10][-8,8)\times[0,10]256×201256\times 201
Burgers (Gaussian initial state) [ 3 ]ut=−uux+0.1 uxxu_{t}=-uu_{x}+0.1\,u_{xx}[−8,8)×[0,10][-8,8)\times[0,10]256×101256\times 101
Burgers (shock) [ 46 ]ut=−uux+(0.01/π) uxxu_{t}=-uu_{x}+(0.01/\pi)\,u_{xx}[−1,1]×[0,0.99][-1,1]\times[0,0.99]256×100256\times 100
2D Burgers [ 10 ]ut=−uux−uuy+0.01 uxx+0.01 uyyu_{t}=-uu_{x}-uu_{y}+0.01\,u_{xx}+0.01\,u_{yy}[−1,1]2×[0,1.98][-1,1]^{2}\times[0,1.98]101×51×100101\times 51\times 100
KdV [ 5 ]ut=−uux−0.0025 uxxxu_{t}=-uu_{x}-0.0025\,u_{xxx}[−1,1)×[0,1][-1,1)\times[0,1]256×201256\times 201
Kuramoto–Sivashinsky [ 47 ]ut=−uux−uxx−uxxxxu_{t}=-uu_{x}-u_{xx}-u_{xxxx}[−10,10)×[0,50][-10,10)\times[0,50]511×251511\times 251
Appendix
Table B2: PDE problems.
More in this topic

AI Agents for Scientific Discovery

AI agents for scientific discovery use language models, tools, and experiments to generate, test, and refine scientific knowledge.

See all 288 papers

Explore similar work

CardsList

  1. STRIDE: A Self-Reflective Agent Framework for Reliable Automatic Equation Discovery

    May 18, 2026Jiarui Su, Songjun Tu, Bei Sun +1LLM Self-RefinementAI Agents for Scientific Discovery

    Central South University · Institute of Automation, Chinese Academy of Sciences · Pengcheng Laboratory

  2. Learning to Design SR Solvers: Harness Optimization for Symbolic Regression

    Oct 12, 2026Qian Li, Yuxiao Hu, Qihao Duan +2Symbolic RegressionAgent Harness Optimization

    Shanghai Jiao Tong University, Shanghai, China · Eastern Institute of Technology, Ningbo, China · The Hong Kong Polytechnic University, Hong Kong, China +1

  3. Language models guide symbolic equation discovery by controlling search

    Jul 5, 2026Zikai Xie, Wenmei Li, Man Luo +2Large Language Model-Guided Scientific DiscoverySymbolic Regression

    State Key Laboratory of Precision and Intelligent Chemistry, Hefei National Research Center for Physical Sciences at the Microscale, School of Chemistry and Materials Science, University of Science and Technology of China, Hefei, China. · Center for Scientific Intelligence Innovation, Hefei, China · School of Chemistry, School of Computer Science, University of Birmingham, Birmingham, UK