Skip to content

Repository files navigation

AgentFold

AgentFold: Closed-Loop Agentic Search for Protein Folding Model Design

AgentFold is a multi-agent framework that searches for executable protein-folding model improvements. Starting from an ESMFold-derived model, it proposes architecture changes, edits and debugs code, trains candidate variants, evaluates predicted structures, and stores both successful and failed experiments for later search.

AgentFold overview

Method

AgentFold runs a closed loop:

  1. Sample: select a parent implementation from the code-variant tree.
  2. Propose and implement: summarize prior evidence, reject duplicate ideas, and edit the folding model.
  3. Debug and train: repair syntax/runtime errors and launch training.
  4. Evaluate and analyze: score structures with OpenStructure and summarize the intervention outcome.
  5. Update memory and search: store code, metrics, failures, and analysis; periodically update MCTS-style priorities.

Each tree node is an executable implementation. The search combines objective training/evaluation signals with a critic score used only for scheduling later experiments.

Main Results

The paper compares the first 36 evaluations from each method under the same model, training budget, and evaluator. Top-5 variants are selected by NWRS.

Method Best lDDT Top-5 lDDT Best NWRS Top-5 NWRS
AgentFold 0.285 0.267 0.526 0.516
Codex proposals 0.265 0.257 0.512 0.509
Random controller 0.260 0.242 0.510 0.506

The strongest reported variant, esmfold_struct_enhanced_v4, improves NWRS by +0.026 and mean/median lDDT by +0.053/+0.059 relative to the fixed baseline.

Repository

agent/
  asi_arch_langchain/       Agent controller and prompts
folding_runtime/
  source/                   Folding and evaluation code
  patches/                  Training/evaluation runtime patch
variants/
  random36_evaluated/       Exact 36 Random-controller variants in Table 1
  codex36_evaluated/        Exact 36 Codex-only variants in Table 1
  random80_success_only/    80 runnable random-control variants
  codex10_baseline/         Original Codex generation artifacts
ops/                        Generation, queue, and analysis utilities
docs/                       Method figure

The folding source also contains the main AgentFold candidate programs and architecture-tree metadata under:

folding_runtime/source/work/alphaflow/programs/
folding_runtime/source/work/alphaflow/program_new/

Environment

The reference environment uses Linux, Python 3.9, CUDA 11, PyTorch 1.12.1, PyTorch Lightning 2.0.4, OpenFold, fair-esm, and OpenStructure.

Create the folding environment from the included specification:

cd folding_runtime/source/work/alphaflow
conda env create -f environment.yml
conda activate alphaflow

The following large artifacts are not included:

Artifact Expected path
ESMFold weights folding_runtime/source/work/alphaflow/esmfold_3B_v1.pt
Training NPZ folding_runtime/source/work/alphaflow/random_npz/
Validation structures folding_runtime/source/work/alphaflow/pdb_test_cif/
Dataset splits folding_runtime/source/work/alphaflow/splits/
Evaluation data folding_runtime/source/work/alpha/data/
Evaluation targets folding_runtime/source/work/alpha/target/

Use these artifacts under their original licenses and terms.

Training

The reported experiments use the patch in folding_runtime/patches/ to disable checkpoints by default, run structural evaluation, save CSV metrics, and remove temporary structure files.

Apply it from the repository root:

cp folding_runtime/patches/rebuttal_no_checkpoint_auto_eval_cleanup/train_custom.py \
  folding_runtime/source/work/alphaflow/train_custom.py

cp folding_runtime/patches/rebuttal_no_checkpoint_auto_eval_cleanup/wrapper.py \
  folding_runtime/source/work/alphaflow/alphaflow/model/wrapper.py

Configure the evaluation entry point:

export AGENTFOLD_EVALUATION_SCRIPT="$(pwd)/folding_runtime/source/work/alpha/evaluation.py"
export AF_SAVE_TOP_K=0

Run one variant:

cd folding_runtime/source/work/alphaflow

./run.sh \
  1 none 410 1000 ./splits/pdb_subset_10k.csv \
  150 5 101 <MODEL_INDEX> <OUTPUT_DIR> 1

The positional arguments are:

trunk_blocks weights max_val_length max_train_samples train_csv
epochs validation_frequency checkpoint_frequency model_index output_dir seed

For a smoke test:

./run.sh \
  1 none 128 8 ./splits/pdb_subset_10k.csv \
  1 1 101 <MODEL_INDEX> smoke_runs/model_<MODEL_INDEX> 1

Evaluation

Run OpenStructure evaluation independently:

python3 folding_runtime/source/work/alpha/evaluation.py \
  --data_dir <PREDICTED_PDB_DIR> \
  --out_dir <ALIGNMENT_OUTPUT_DIR> \
  --dataset cameo22

The evaluator reports backbone lDDT, lDDT, oligomeric GDT-TS, RMSD, and TM-score. The included analysis utilities summarize mean and median metrics and select the best evaluated epoch by mean backbone lDDT.

Agent Search

Set an LLM API key through the environment:

export DASHSCOPE_API_KEY=<YOUR_KEY>
# or
export OPENAI_API_KEY=<YOUR_KEY>

The main entry points are:

agent/asi_arch_langchain/main.py
agent/asi_arch_langchain/evolve_file.py
agent/asi_arch_langchain/run_alphaflow_random80.sh

Do not store API keys in source files or commit a .env file.

Citation

@article{liu2026agentfold,
  title   = {AgentFold: Closed-Loop Agentic Search for Protein Folding Model Design},
  author  = {Liu, Mingquan and Chen, Jiangyu and Cao, Hanqun and Zhang, Xujun and
             Ma, Pengsen and Tang, Xiangru and Jin, Shuting and Yang, Zhuo and
             Fu, Tianfan and Wu, Fang and Zeng, Xiangxiang},
  year    = {2026},
  note    = {Preprint}
}

License

The repository includes components under different upstream licenses. See THIRD_PARTY_NOTICES.md before redistribution.

About

No description, website, or topics provided.

Resources

Stars

10 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages