Reasoning-driven Agentic Search across prompts, programs, and ML workflows.
Junbo Li1,
Boyi Liu2,
Canwen Xu2,
Yite Wang2,
Zhewei Yao2,
Zhangyang Wang1,
Qiang Liu1,
Yuxiong He2
1UT Austin 2Snowflake AI Research
Prior LLM-based optimizers slot the model into an external scaffold — bandits, evolutionary loops, textual gradients — that decides which candidate to branch from, when to verify, how to allocate budget, and when to give up. ReASearch removes the scaffold. A single tool-using reasoning agent owns the whole search policy: it analyzes outcomes, allocates its own budget, deliberately reverts unproductive branches, double-verifies promising gains, and adapts exploration as it goes. We instantiate the same agent on three substrates — prompts (e.g., AIME, GSM8K, HotpotQA), programs (e.g., Circle Packing, ARC-AGI), and ML training pipelines (e.g., image classification, NanoGPT, Atari/MuJoCo RL) — and across 14 tasks it consistently matches or outperforms specialized systems built around explicit search procedures.
Beyond results, the framework is built to be easy to extend and effectively controllable. Hooking a new optimization domain in is a handful of methods on a base class — give the agent a metric and an editable artifact (prompt, function, or training script), and the same reasoning loop, persistent memory, and search-review machinery just work (see Add your own task). Every lever on the search itself is exposed: budget caps (--max-total-evaluations, --max-agent-turns, --max-wall-clock), declarative constraints (no_change for forbidden moves, with checker-LLM enforcement), seeded prior knowledge (--lessons), and explicit search-tree review — so you can shape how the agent searches, not just whether it runs.
- Architecture
- Quick start
- Task families
- Selected results
- Key features
- Add your own task
- LLM providers
- Outputs
- Documentation
ReASearch/
├── run.py # entry point — dispatches to task.run()
├── config.py # global config (providers, budgets, paths)
│
├── core/ # task-agnostic agent runtime
│ ├── agent.py # the reasoning loop (one function, no controller)
│ ├── context.py # auto-compaction; lessons.md survives
│ ├── tools.py # python_exec / bash / read_file / write_file / list_files / compact
│ ├── state.py # AgentState container
│ ├── prompts.py # shared prompt utilities
│ └── opencode.py # opencode CLI bridge for repo-mode editing
│
├── clients/ # six LLM provider adapters with a unified message+tool interface
│ ├── openai_client.py # OpenAI
│ ├── anthropic_client.py # Anthropic
│ ├── gemini_client.py # Google Gemini (OpenAI-compat)
│ ├── azure_openai.py # Azure OpenAI
│ ├── bedrock.py # AWS Bedrock
│ └── cortex.py # Snowflake Cortex
│
└── tasks/
├── base.py # BaseTask, search-review, shared helpers
├── eval_store.py # evaluation caching (deduplication)
├── prompt_optimization/ # aime · gsm8k · hotpotqa · terminal_bench
├── program_evolution/ # arc_agi · circle_packing · heilbronn · eplb · txn_scheduling
└── ml_workflow_optimization/ # nanogpt · market · atari · mujoco
Prerequisites: Python 3.10+, uv (ML-workflow tasks use it for per-task Python envs), an LLM provider key (see LLM providers), and a CUDA GPU for ML-workflow and some program-evolution tasks.
# 1. Clone and install
git clone <repo-url>
cd ReASearch
pip install -r requirements.txt
# 2. Configure a provider
cp .env_template .env
# edit .env: set OPENAI_API_KEY=... (or ANTHROPIC_API_KEY=...) and AGENT_PROVIDER=openai
# 3. Run a task
python run.py \
--task nanogpt \
--max-agent-turns 1000 \
--enable-search-review \
--exp-name my_first_runOutputs land in workspace/<task>/<timestamp>_<exp-name>/. Every prompt edit, every experiment, every checker decision is captured.
Optional — opencode for --repo-mode. With --repo-mode, an ML-workflow task edits a whole multi-file repo via opencode. Install with curl -fsSL https://opencode.ai/install | bash; it reads the same .env credentials you already set. Snowflake Cortex is the one provider that needs an extra step — run bash scripts/setup_opencode.sh once. See OPENCODE_SETUP.md.
| Family | What the agent edits | Evaluator | Tasks |
|---|---|---|---|
| Prompt Optimization | A system or user prompt | Student LLM on train/val/test split | aime, gsm8k, hotpotqa, terminal_bench |
| Program Evolution | A solve() function |
Direct execution against a problem instance | arc_agi, circle_packing, heilbronn, eplb, txn_scheduling |
| ML Workflow Optimization | A training script (single-file or full repo) | uv run train.py reporting a metric |
nanogpt, market, atari, mujoco |
Each family has its own README with task-specific setup notes. The same agent loop drives all three.
| Method | AIME 2025 | GSM8K | HotpotQA | Terminal-Bench 2.0 |
|---|---|---|---|---|
| Baseline | 46.00 | 81.20 | 63.00 | 35.56 |
| GEPA | 50.67 | 82.11 | 65.80 | 42.20 |
| ReASearch | 52.00 | 83.40 | 67.60 | 53.33 |
Here is an example. The agent grew the HotpotQA prompt from a one-line stub into a four-step CoT protocol with answer-type rules, span-extraction discipline, and a verification step:
| Baseline (V1, test 63.0%) | ReASearch best (V14, test 67.6%) |
|---|---|
| You are a question answering assistant. You will be given context documents and a question. Read the documents carefully and answer the question. Output your final answer as ### <answer> on its own line. |
You are an expert question answering assistant for multi-hop reasoning tasks… Step 1 — Build the reasoning chain: trace fact A → fact B → final answer. Step 2 — Determine the answer type from the question: Who… → person; Which entity… → that entity; How many… → number; Is/Are/Was… → yes/no only. Step 3 — Extract the answer span: a direct quote, minimum phrase, no surrounding context. Step 4 — Verify your answer: right type? right entity? not too much, not too little? |
Circle packing (sum of radii in unit square, Sonnet 4.6):
| n | Human best | AdaEvolve | ReASearch |
|---|---|---|---|
| 26 | 2.635 | 2.630 | 2.636 |
| 30 | 2.842 | 2.843 | 2.843 |
| 31 | 2.889 | 2.887 | 2.890 |
| 32 | 2.939 | 2.810 | 2.940 |
Heilbronn triangle (minimum triangle area, Sonnet 4.6):
| n | Human best | AdaEvolve | ReASearch |
|---|---|---|---|
| 12 | 0.03260 | 0.02552 | 0.03260 |
| 13 | 0.02702 | 0.02142 | 0.02700 |
| 14 | 0.02430 | 0.00299 | 0.02429 |
ARC-AGI-2 (120 puzzles, Sonnet 4.6, pass@2):
| Train acc | Test acc | |
|---|---|---|
| AdaEvolve | 21.9% | 12.5% |
| ReASearch | 85.0% | 50.0% |
Systems programming (best-of-paper, GPT-5 backbones; TXN: makespan score ↑; EPLB: balance × speed ↑):
| TXN | EPLB | |
|---|---|---|
| AdaEvolve | 3,636 | 0.1976 |
| ReASearch | 4,237 | 0.2305 |
 |
| Task | Direction | Train budget | Baseline | ReASearch | Discovered configuration |
|---|---|---|---|---|---|
| Image classification (from scratch) | val acc | 5 min | 58.7 | 83.3 | ResNet-34 widened [96,192,384,512] + AutoAugment + SWA + TTA + bs=512; the +14-pt jump came from the agent computing per-step wall-clock and realizing the cosine LR schedule never reached its low-LR phase in 5 min (estimated_epochs=100 fix) |
| Atari Q*bert (PPO) | reward | 30 min | 475 | 18,600 | 48 envs · 8 PPO epochs · entropy=0.001 · 12 minibatches; the 5× jump came from deepening FC layers (512→1024→512) after the agent observed that doubling envs from 8→16 left reward unchanged — learning efficiency, not data quantity, was the bottleneck |
| MuJoCo HalfCheetah (TD3) | reward | 10 min | 3,947 | 7,840 | TD3 bs=384, min(Q1,Q2) clipping, buffer=500k, reward_scale=5.2; agent benchmarked priority-sampling vs uniform np.random.randint, measured a 133× slowdown, and rejected PER on computational grounds before wasting an experiment |
| NanoGPT (FineWeb-Edu) | val_bpb | 5 min | 0.9929 | 0.9696 | Muon momentum ramp + softcap=15 + RoPE/warmdown/weight-decay co-tuning, after the agent diagnosed batch-size co-dependence (lr 0.04→0.027 when bs goes 2¹⁸→2¹⁷) |
| Crypto (DRW Kaggle) | val_corr (IC) | 15 min | 0.0953 | 0.1110 | Ridge + MLP ensemble + per-feature winsorization (2.8–97.2 pctl) under purged time-series CV; within 15 experiments the same configuration moved the Kaggle leaderboard from 36th → 6th |
The framework is small. The interesting part is the cognitive scaffolding around the agent — features that let the model itself drive search rather than be driven.
The agent isn't just emitting candidates — it has a real workspace and a Python REPL. python_exec runs in a persistent namespace (variables, loaded dataframes, model handles all survive across turns), with json / re / math / statistics / collections pre-loaded and the cwd set to the run dir so open("logs/evals/...") just works. Pair that with workspace-scoped read_file / write_file / list_files and a sandboxed bash (120s timeout, dangerous-command deny-list), and the agent becomes a capable analyst — it inspects its own training logs, plots score histories, greps eval outputs, and writes helper scripts while it's deciding what to try next.
Every run starts with an empty lessons.md in the workspace. The agent is instructed to read it at the start of each session and update it as it learns: what architectures or hyperparameters worked, what didn't, what to try next. Critically, the file is re-injected into the system prompt on every turn (core/context.py), so when the conversation gets compacted at long horizons, lessons survive. Without this the agent forgets ablation results around turn 200 and starts repeating itself.
Pass --lessons PATH to resume from prior memory — the file is copied into the new run's workspace at startup, so a follow-up run picks up where a previous one left off (or starts from a hand-curated set of takeaways). Useful for chaining long horizons, swapping models mid-search, or transferring insights from a sibling task.
Adds a review_search_tree tool. When the agent calls it, the tool renders the experiment history as a structured tree with branch concentration, score progression, plateau signals, and a four-question analysis prompt (plateau? diversity? budget allocation? next action?). The agent typically calls it every 3-5 experiments. In our runs, the output drove real strategy decisions — e.g., "The search tree shows we're on the right track with ResNet-18 + AMP. The key insight is that larger models all fail because they run fewer epochs in 5 minutes."
### Search Tree
#attempt_0001 score=60.93 Baseline: ResNet-18, SGD, cosine LR
├── #attempt_0002 score=45.85 [discard] WideResNet-28-10 + AutoAugment + CutMix
└── #attempt_0003 score=63.75 ★ ResNet-18 + TrivialAugmentWide + CutMix + AMP
├── #attempt_0004 score=55.37 [discard] ResNet-50 + TrivialAugment + CutMix
└── #attempt_0005 score=60.83 [discard] WideResNet-28-4 + TrivialAugment
tasks/program_evolution/prompts.py injects a ## Strategy Advisory section into the system prompt when no improvement has occurred for 3+ evaluations. After 7 evaluations of stagnation it escalates: "Your current strategy isn't working — try something fundamentally different." Each task can also override _stagnation_suggestions() to add domain-specific kick-out advice. This is the difference between an agent that grinds productively and one that loops on local optima.
A task can declare a no_change property — a free-form natural-language constraint the agent must respect. The constraint is shown to the agent in its system prompt and every code edit is reviewed by a checker LLM that compares the diff against the rule. Violating diffs are auto-reverted with feedback to the agent; decisions log to logs/no_change_checks.jsonl. Concrete uses:
- Image classification (from scratch): forbid pre-trained weights / transfer learning.
market: ban temporal leakage in time-series feature engineering.
In practice the rule shapes search up-front: the agent reads the constraint section and plans around it, rather than discovering it via repeated reverts.
For ML-workflow tasks that ship multi-file repos rather than a single train.py, the agent uses understand_repo and edit_repo (powered by opencode) instead of edit_train_file. The same multi-file pattern is also a natural fit for program-evolution problems whose solution is structurally complex — for example, evolving a search algorithm whose state, heuristic, and scoring components are easier to reason about as separate modules, or a multi-stage solver (preprocess → core → post-process) where the agent benefits from editing each stage independently rather than threading every change through a single solve() function. (Opencode reads its credentials directly; for OpenAI / Anthropic / Gemini / Bedrock / Azure this is the same env keys you've already set. Snowflake Cortex is the one provider that needs an extra step — see OPENCODE_SETUP.md.)
core/context.py automatically compresses the conversation when context grows large. state_summary() provides the authoritative state — current experiment count, budget remaining, recent score progression, and the best candidate — that survives compression. The agent never wakes up amnesic.
Hooking a new domain into ReASearch is the main supported extension path. The agent loop is task-agnostic; you only specify a few methods.
| Family | Base class | Core methods you implement |
|---|---|---|
| Prompt Optimization | PromptOptTask |
load_dataset(), default_prompt(), extract_answer(), score_answer() |
| Program Evolution | ProgramEvoTask |
load_problem(), default_code(), evaluate(code, problem, timeout) |
| ML Workflow | MLWorkflowTask |
metric_name, metric_direction, template_files(), verify_data() |
Plus a system_prompt_task_section() that gives the agent task-specific guidance, and a one-line registration in tasks/__init__.py. You can optionally override no_change to gate edits, _stagnation_suggestions() for domain-specific stuck advice, or post_run_evaluation() for held-out scoring. SKILLS.md is written as a brief that an AI coding agent (Claude Code, Codex, etc.) can read end-to-end and use to scaffold a new task in this repo without further help — point your coding assistant at it.
Set AGENT_PROVIDER in .env:
| Provider | Value | Auth | opencode (repo-mode) |
|---|---|---|---|
| OpenAI | openai |
OPENAI_API_KEY |
works out of the box |
| Anthropic | anthropic |
ANTHROPIC_API_KEY |
works out of the box |
| Google Gemini | gemini |
GEMINI_API_KEY |
works out of the box (via OpenAI-compat endpoint) |
| Azure OpenAI | azure_openai |
AZURE_ENDPOINT + AZURE_API_KEY |
works out of the box |
| AWS Bedrock | aws_bedrock |
AWS creds + AWS_REGION |
works with Bedrock invoke perms |
| Snowflake Cortex | cortex |
SNOWFLAKE_PAT |
needs scripts/setup_opencode.sh |
See .env_template for all configuration options.
| Flag | Description |
|---|---|
--task NAME |
Which task to run |
--max-total-evaluations N |
Evaluation budget (default 500) |
--max-agent-turns N |
Maximum agent reasoning turns |
--max-wall-clock SECONDS |
Wall-clock limit |
--lessons PATH |
Seed agent with an existing lessons.md |
--enable-search-review |
Give the agent the review_search_tree tool |
--repo-mode |
Multi-file editing via opencode (ML-workflow or program-evolution tasks with a repo template) |
--exp-name NAME |
Tag the run directory |
--agent-model-id MODEL |
Override agent model |
Each task may add its own flags; run python run.py --task <name> --help to see them.
Every run produces a workspace directory:
workspace/<task>/<timestamp>_<exp-name>/
├── repo/ # task working tree (where edits land)
├── experiments.json # full kept/discarded experiment history
├── lessons.md # agent-curated insights
├── search_review.md # latest search-review snapshot (if enabled)
└── logs/
├── agent.log # reasoning trace
├── tool_calls.jsonl # one JSON per tool call
├── turn_log.md # human-readable turn-by-turn summary
├── code_snapshots/ # before/after diffs of each edit
├── runs/exp###_run.log # per-experiment train.py output
├── opencode_runs/ # (repo-mode only) opencode subprocess logs
└── no_change_checks.jsonl # (constrained tasks only) checker decisions
| Doc | Audience |
|---|---|
README.md |
Project overview (you are here) |
SKILLS.md |
Brief for an AI coding agent adding a new task to this repo |
OPENCODE_SETUP.md |
opencode auth setup (Cortex-specific path) |
tasks/*/README.md |
Per-task-family setup notes |
If you use this framework in research, please cite:
Li, J., Liu, B., Xu, C., Wang, Y., Yao, Z., Wang, Z., Liu, Q., He, Y. The Optimizer Is the Agent: Reasoning-Driven Search across Prompts, Programs, and ML Workflows. Preprint, 2026.