MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model

Brian Ko; Han Wang; Hejia Geng; Li Erran Li; Sherry Yang

cs.LGOct 27, 2026

MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model

Authors: Brian Ko, Han Wang, Hejia Geng, Li Erran Li, Sherry Yang

Organizations: University of Washington · University of Illinois at Urbana-Champaign · University of Oxford · Columbia University ·

Abstract

AI research agents can propose, execute, and improve programs that solve machine learning tasks. However, executing each candidate program requires training and evaluating ML models, limiting how many candidates an agent can explore at each step. We introduce MLE-World, a world model finetuned from Qwen3.5-4B that predicts a program's terminal output (error traces, training logs, and validation scores) from its code and data description without executing it, together with a confidence score estimating the prediction's correctness. Because a prediction costs one inference call rather than a training run, MLE-World makes test time scaling (TTS) affordable for these agents. We build five TTS policies, guided by MLE-World, that compare several candidates per step by their predicted outputs: Sequential revision, Best-of-N, Pairwise judging, Fusion, and a UCB style rule that also adds an exploration bonus from this confidence score. Since predictions can be wrong, we run a candidate for real only when its confidence score is low. On 17 MLE-bench Lite tasks, MLE-World matches the real terminal output on 77.9% of held out programs, versus 41.1% for much larger prompted models. MLE-World guided TTS achieves up to 22.1% higher observed mean performance than the agent without a world model, with less real execution. Project webpage: https://mle-world.github.io/.

Figures & tables

Figure 2: Five ways to search with MLE-World. Each tile pairs a candidate program with its predicted execution output. Sequential revises programs using predicted feedback; Best-of-

N

selects the candidate with the best predicted score; Pairwise compares candidates in pairs; Fusion combines ideas into a new program; and UCB balances predicted score with an exploration bonus for uncertain candidates. Each policy uses four candidates and returns one program.

Figure 5: The exploration bonus yields a higher observed mean. All other settings are fixed; only Improve-stage

\beta

varies (Figure 8 a).

Figure 7: Prediction fidelity rises with confidence. Accuracy by

p_{\max}

quintile on 1,700 held out programs: (a) metric within 5% relative error, (b) exact exception class, (c) judged terminal output. Bars:

\pm

task bootstrap SE.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

	Programs	%
All training programs	497,056	100.0
Generator model
Qwen3-Coder-Next	209,573	42.2
Qwen3.5-4B	129,907	26.1
Qwen3.5-122B-A10B	77,239	15.5
Qwen3.5-9B	54,815	11.0

Appendix

Table 5: Training data statistics. Most training programs fail: 52.3% crash and 7.0% time out, so MLE-World sees failures more often than clean runs. Counts over the 497,056 training programs, each recovered to its task, generator model, and stage from the deduplicated node pool. Outcome, exception class, and labels come from the training target; token counts use the Qwen3.5-4B tokenizer, and the prompt includes the system message.

Base model	Qwen3.5-4B
Method	SFT of all parameters (LlamaFactory), vision tower frozen
Training data	497,056 programs (Table 5 )
Chat template	qwen3_5 , thinking disabled
Added tokens	82 exception class tokens ( EXC_None , EXC_ValueError , …; Appendix C )
Sequence length	12,288 tokens, packed without attention across examples
Epochs / steps	5 / 6,685

Appendix

Table 6: Training configuration. SFT of all parameters of Qwen3.5-4B; we evaluate the final checkpoint (step 6,685).

None ( null )	MissingCorpusError
AcceleratorError	ModuleNotFoundError
ArrowInvalid	NameError
ArrowTypeError	NotADirectoryError
AssertionError	NotFittedError
AttributeError	NotImplementedError
AxisError	OSError

Appendix

Table 7: Exception classes. The 82 classes of

\mathcal{C}

, in token order; each is one added token EXC_ name .

target	Terminal output	Exception class	Crash class	Validation metric (precision / recall)
terminal output (untrimmed) + submission flag	59.5	71.4	64.4	87.3 / 83.5
terminal output (trimmed) + submission flag	59.4	71.7	64.8	88.7 / 84.3
+ exception class	60.7	73.0	66.7	87.3 / 87.6
+ validation metric	62.8	75.3	69.1	89.0 / 89.5
validation metric + submission flag	–	–	–	90.0 / 77.0

Appendix

Table 8: Target ablation. Structured fields improve the predicted log, and dropping the log costs metric recall. 450 held out programs from its nine training tasks; backbone, data, and decoding are fixed. “+” adds fields cumulatively; trimmed logs retain the first and last 2,500 characters. All scores are percentages. Terminal output measures semantic agreement; exception/crash measure exact class agreement inferred from logs, with “no exception” included and crash restricted to jointly crashing cases. Metric precision uses 10% relative tolerance on commonly answered cases; recall is the fraction of cases with a numeric reference receiving a metric. Dashes indicate undefined scores from logs. Each row is one training run; at 450 programs, a rate near 60–75% has a bootstrap SE over programs of about 2 points (more in the crash and metric columns, which use subsets).

	MLE-World	Coder	DeepSeek	Judge (Coder)
Reasoning	off	off	on (max effort)	off
Temperature	0 (greedy)	0 (greedy)	1.0	1.0
Top- $p$	1.0	1.0	0.95	0.95
Top- $k$	–	–	–	40
Samples per program	1	1	1	3
Max output tokens	8,192	8,192	32,768	3,072

Appendix

Table 9: Decoding settings. Settings used to predict the 1,700 held out programs, and for the judge that scores terminal output and crash class.

Method	Terminal output	Exception class	Crash class	Validation metric
Task majority	–	34.2 $\,\pm$ 1.2	–	55.2 $\,\pm$ 2.3
Nearest program	28.0 $\,\pm$ 1.1	41.1 $\,\pm$ 1.2	14.0 $\,\pm$ 1.0	55.6 $\,\pm$ 2.2
Parent record	22.3 $\,\pm$ 1.0	34.2 $\,\pm$ 1.2	6.7 $\,\pm$ 0.7	48.8 $\,\pm$ 2.3
Random record	11.2 $\,\pm$ 0.8	22.2 $\,\pm$ 1.0	3.5 $\,\pm$ 0.6	20.2 $\,\pm$ 1.8
MLE-World	77.9 $\,\pm$ 1.0	86.4 $\,\pm$ 0.8	76.4 $\,\pm$ 1.3	78.1 $\,\pm$ 1.9

Appendix

Table 10: Corpus baselines. Baselines built from MLE-World’s training corpus without a model trail MLE-World on every field and stage. Accuracy (%,

\pm

SE from 2,000 bootstrap resamples over programs) on the 1,700 held out programs; stage columns give Draft/Improve/Debug. Parent record equals Nearest program on the 561 drafts; their draft terminal output scores differ by 0.1 only because each baseline was judged separately by the sampled judge.

	DeepSeek	Coder	MLE-World (4B + SFT)
Finished before token cap	94.9	91.8	99.7
Exception class	93.6	100.0	100.0
Terminal output	91.8	91.8	99.7
Validation metric	91.8	91.8	99.7
All three fields	91.4	91.8	99.7
Median output tokens	12,042	144	160

Appendix

Table 11: Output parse rates. Prompted models rarely fail to parse, so parsing does not explain their low accuracy. Share (%) of the 1,700 held out predictions that finish before the token cap and contain each scored field; a missing exception field is read from the terminal output. Both prompted models return all three fields for over 91% of programs, far too few failures to account for their 28–60 percentage point gap in Table 1 . DeepSeek’s output is long because it reasons, while the other two models run with reasoning off.

Stage
		Terminal output			Exception class
	$n$	MLE	DS	Coder	MLE	DS	Coder
Draft	561	78.3	40.9	16.5	86.5	55.4	26.6
Improve	561	83.2	38.9	21.4	89.8	57.2	29.8
Debug	578	72.3	43.5	15.5	82.9	59.0	26.3
Outcome

Appendix

Table 12: Accuracy by stage, outcome, and crash type. MLE-World leads both prompted baselines on every slice and field except the crash class of IndexError (26 programs), and timeouts are its hardest case. Accuracy (%) of MLE-World (MLE), DeepSeek (DS), and Coder;

n

counts programs per row. Darker cells are more accurate and the best value per field is bold. Crash type splits the 1,094 crashing programs by their observed exception class; the ten most frequent are listed. Crash class is scored only on program errors and the metric only on programs with a numeric reference, so some cells are empty.

	Cases	Agreement
All cases	120	96.7
Actual failures	82	95.1
Crashes of the same class	40	93.3
Clean runs	20	100.0

Appendix

Table 13: Judge validation on the held out test set. Majority agreement of the final judge prompt with the frozen human labels (%).

Improve stage, eight hours; real check after each selection
Selection rule	Mean percentile	Valid submissions
Metric greedy, $\beta=0$	38.58 $\pm$ 3.52	38/51
Metric UCB, $\beta=1$	45.38 $\pm$ 2.36	44/51
Debug stage, eight hours; real check after each selection
Greedy, $p_{\mathrm{clean}}$	41.88 $\pm$ 2.10	48/51
UCB, $v+\beta u$ , $\beta=1$	48.37 $\pm$ 1.82	48/51

Appendix

Table 14: UCB ablations by stage. Compare rows within each block: Improve isolates the bonus; Debug changes reward and eligibility. Mean percentile

\pm

SE across three seed means.

Figure 8: Search trajectories for the ablations by stage. (a) Improve: metric greedy (

\beta=0

) versus UCB (

\beta=1

). (b) Debug:

p_{\mathrm{clean}}

greedy versus metric UCB. Green: greedy; blue: UCB. Curves carry forward the latest reviewed, validation-selected submission, scoring zero before the first valid submission. Endpoints match independent final grades.

(a) Ungated comparison at three hours
Setting	$\beta$	Mean percentile	Valid submissions
Metric only BoN reference †	0	36.47 $\pm$ 0.10	38/51
$p_{\mathrm{clean}}$ UCB	0.25	37.28 $\pm$ 4.87	38/51
$p_{\mathrm{clean}}$ UCB	0.5	41.71 $\pm$ 2.89	38/51
$p_{\mathrm{clean}}$ UCB †	1	38.07 $\pm$ 1.93	41/51
$p_{\mathrm{clean}}$ UCB	2	34.82 $\pm$ 0.92	35/51

Appendix

Table 15: Exploration weight and confidence gating. Mean leaderboard percentile

\pm

SE across the three seed means; every row retains all 51 trials. These 8 h controls precede the 12 h all-stage UCB results in Tables 2 and 4 . Comparisons across experiment waves are descriptive.

Figure 9: Holdout prediction error versus skipped execution time. Raw

p_{\max}

gates on eligible Improve and Debug programs. Markers show the two success thresholds (0.90, 0.50) and two failure thresholds (0.83, 0.56). Success markers use the criterion of relative metric error exceeding 5%. Failure markers count exception or cause mismatches. Runtime percentages are computed separately within each predicted outcome. The held out programs belong to the MLE-World evaluation of Section 6.1 , not to any agent run. These estimates do not fully reflect the

p_{\max}

distribution each policy produces during search; they propose candidate thresholds, and a separate development run chooses among them (Table 16 ).

Policy	Ungated	$(0.90,0.83)$	$(0.50,0.83)$	$(0.50,0.56)$
	Mean percentile (mean iterations)
		Gated $(\gamma_{\mathrm{s}},\gamma_{\mathrm{f}})$
Sequential	32.31 (22.0)	37.39 (27.0)	20.86 (30.5)	47.20 (37.6)
BoN	36.47 (29.6)	44.90 (49.1)	48.30 (69.4)	52.25 (58.0)
Pairwise	40.44 (28.5)	40.19 (44.1)	38.60 (44.7)	45.26 (71.6)
Fusion	34.01 (24.9)	38.90 (51.9)	44.59 (44.6)	39.48 (45.4)

Appendix

Table 16: Gate threshold sweep at three hours. Each cell gives mean leaderboard percentile. Every policy’s best pair uses

\gamma_{\mathrm{s}}=0.50

	Random $p_{\mathrm{run}}=0.50$			Confidence gate		$\Delta$ percentile
Policy	Valid	Pct.	Hours	Pct.	Hours	$\pm$ SE
BoN	46/51	47.37	5.54	47.21	4.49	$-0.15\pm 2.32$

Appendix

Table 17: BoN: random execution versus confidence gating after eight hours. Each gate has 51 trials. Percentile includes invalid submissions as zero; hours are real execution per trial.

\Delta

is confidence minus random, paired by task.

Figure 10: Gated mean percentile over time. Five confidence-gated policies from Table 4 and Base AIDE; 17 tasks

\times

three seeds, 12 h per trial. Curves carry forward the validation-selected submission, with zero before the first valid submission. Comparisons across archived cohorts are descriptive.

Evaluation point	Mean percentile $\uparrow$	$\Delta$ vs. Base	Valid (%) $\uparrow$
8 h cutoff	41.9 $\pm$ 3.5	+5.1 $\pm$ 3.4	88.2 $\pm$ 3.4
Final (16 h cap)	45.0 $\pm$ 2.3	+8.1 $\pm$ 6.2	94.1 $\pm$ 3.4

Appendix

Table 18: Real-execution BoN reference. 17 tasks

\times

three seeds; all four candidates execute at each step. Mean percentile and valid submission rate:

\pm

SE across seed means;

\Delta

versus the earlier 8 h Base AIDE result:

\pm

paired task bootstrap SE. The final row uses a larger time cap and is not budget matched to Base AIDE. Invalid submissions score zero.

Policy	Valid (%)	Mean pct.	Pct. if valid	First (h)	Trials with one	Per trial
	Final submission			Verified new bests
Base AIDE	78.4	36.9	47.0	0.43	40	3.4
MLE-World guided
Sequential	90.2	41.1	45.5	0.62	46	4.5
+ confidence gate	96.1	47.7	49.7	0.68	49	4.4
Pairwise	92.2	42.8	46.4	1.06	47	5.1

Appendix

Table 19: Agent execution statistics. Guidance raises the share of real runs that finish cleanly, and the confidence gate cuts real runs per trial by 48–73%. The clean run share rises with gating for four of five policies. 51 trials per row (17 tasks, three seeds, eight hours). Final submission: rate of valid submissions, mean leaderboard percentile in these earlier cohorts, and mean percentile over valid submissions only. Real runs: executions per trial and their outcomes, pooled over all runs of a policy (clean: finished without an exception; timeout: hit the execution limit), and hours of real execution per trial, in total and in timed out runs. Verified new bests: median hours to the first verified submission among trials that have one, the number of such trials, and verified new bests per trial.

Selection: pick one of the same four candidates
Rule	Mean percentile	$\Delta$ mean pct. (points)	Valid (%)	$\Delta$ valid (points)
Uniform over the pool	10.9	–	26.1	–
Screen by predicted outcome	16.4	+5.51 $\pm$ 1.18	41.9	+15.8 $\pm$ 3.1
+ ranking by predicted metric	17.1	+0.68 $\pm$ 0.79	42.5	+0.5 $\pm$ 1.3
Real validation of all four †	19.8	+2.65 $\pm$ 2.19	47.1	+4.6 $\pm$ 2.6
Pairwise: judge the same four candidates

Appendix

Table 20: Diagnostics on fixed states. Screening and Fusion improve validity and test percentile; the test percentile gains of Pairwise and Sequential from predicted terminal output are not established. Selection, Pairwise, and Fusion share their input pool; Sequential shares its starting candidate. Uniform : expected outcome of a random pick. Screen : random pick among candidates MLE-World predicts usable with a finite metric, else the largest

p_{\mathrm{clean}}(x_{s})

. + ranking : the screened candidate with the best predicted metric. Real validation : best real validation score, else the ranking’s pick. Sequential uses the same selector by predicted metric and fallback in every condition; Structured predictions omit terminal output from revision feedback. Test pct. is zero if invalid.

\Delta

: paired difference from the row above within each panel

\pm

SE over 4,000 hierarchical task/run bootstrap resamples. † Executes all four inputs. ‡ Uses real execution feedback during revision.

Task	Teams	Gold rank	Gold pct.	Golds	Any medal	Policy pct. (range)
Tasks with a gold
detecting-insults-in-social-commentary	50	5	90.0	17	17	81.3–100.0
dogs-vs-cats-redux-kernels-edition	1,315	12	99.1	4	9	33.3–95.5
plant-pathology-2020-fgvc7	1,318	12	99.1	4	5	10.1–92.9
tabular-playground-series-dec-2021	1,189	12	98.7	4	4	46.3–72.7
aerial-cactus-identification	1,221	12	87.7 †	2	2	50.4–69.8

Appendix

Table 21: Medals by task. More than half of all golds come from one task and twelve tasks yield none, yet policies still differ in percentile on tasks without medals. The 306 final submissions of the six earlier 8 h ungated policies (18 per task: six policies, three seeds). Teams: leaderboard size. Gold rank: the rank a submission must reach for gold. Gold pct.: percentile of a submission scoring exactly the gold threshold. Golds and Any medal count submissions. Policy pct.: lowest and highest of the six mean percentiles of the policies. † Gold requires the best leaderboard score, which many teams share: AUC 1.0 on aerial-cactus (150 teams) and log loss 0.0 on leaf-classification (24 teams).

Policy	detecting- insults	Other	Any medal	Invalid	Gold trials	Other valid trials	Total
	Golds				Mean percentile
Base AIDE	3	3	11	11	11.8	25.1	36.9
Sequential	3	3	11	5	11.5	29.5	41.1
Pairwise	3	5	14	4	15.4	27.4	42.8
Fusion	2	1	11	2	5.9	41.3	47.2
BoN	3	0	9	6	5.9	33.6	39.5

Appendix

Table 22: Medals and percentile by policy. Gold counts turn on a few submissions, whereas most of each policy’s mean percentile comes from submissions without a gold. Counts over 51 trials per policy. Mean percentile (Total, as in Table 2 ) splits into points from gold medal trials and from other valid trials; invalid submissions contribute zero, and rounded components may not sum to the total.

Method	Terminal output	Exception class	Crash class	Validation metric
DeepSeek	42.1 $\pm$ 2.8	63.7 $\pm$ 2.8	45.4 $\pm$ 3.2	21.3 $\pm$ 5.8
Coder	21.1 $\pm$ 2.2	30.3 $\pm$ 2.6	18.9 $\pm$ 2.4	31.9 $\pm$ 6.8
MLE-World (4B, no CoT)	53.3 $\pm$ 2.8	60.3 $\pm$ 2.8	56.3 $\pm$ 3.1	44.7 $\pm$ 7.3

Appendix

Table 23: Exploratory prediction on three tasks excluded from finetuning. Accuracy (%,

\pm

program bootstrap SE) on 300 programs, 100 per task with balanced agent stages. The fields and scoring match Table 1 ; DeepSeek uses default reasoning effort here. Bold denotes the best value in each column.

Held out task	Related training tasks	Shared structure and remaining shift
APTOS Blindness Detection	Plant Pathology 2020; Dog Breed	Image classification with pretrained visual encoders; retinal severity introduces ordinal labels and a new metric.
3D Object Detection	Dog Breed; Aerial Cactus	Shared image, array, and PyTorch components can expose familiar import and file errors; 3D geometry and detection remain new task specific computations.
RANZCR Catheter/Line	Plant Pathology 2020	Image preprocessing and visual prediction heads are shared; chest radiographs and multiple catheter/line labels change the target.

Appendix

Table 24: Shared operations that could support transfer. The middle column names tasks present in finetuning. These are proposed mechanisms, not measured causal effects.

APTOS Blindness Detection
Method	Terminal output	Exception class	Crash class	Validation metric
DeepSeek	49.3 $\pm$ 4.9	68.0 $\pm$ 4.5	59.0 $\pm$ 5.6	16.7 $\pm$ 8.6
Coder	18.3 $\pm$ 3.8	32.0 $\pm$ 4.6	15.3 $\pm$ 4.2	16.7 $\pm$ 9.1
MLE-World (4B, no CoT)	53.0 $\pm$ 4.9	60.0 $\pm$ 4.8	53.6 $\pm$ 5.6	50.0 $\pm$ 12.1
3D Object Detection
DeepSeek	35.7 $\pm$ 4.7	52.0 $\pm$ 4.9	26.5 $\pm$ 5.0	31.8 $\pm$ 10.1

Appendix

Table 25: Per task prediction accuracy on the three selected tasks. Accuracy (%,

\pm

program bootstrap SE) using the same methods and metrics as Table 23 . Each task has 100 programs: 34 Draft, 33 Improve, and 33 Debug. Gray

\dagger

cells have fewer than ten metric references; SE is omitted when the metric is undefined in any bootstrap resample.

Test-Time Scaling

Test-time scaling studies how extra inference-time computation can improve model answers, decisions, and generated outputs.

See all 223 papers

Qwen-AgentWorld: Language World Models for General Agents
Jun 23, 2026Yuxin Zuo, Zikai Xiao, Li Sheng +30LLM World Models Agentic RL
Qwen Team
TuiML: Machine Learning for AI Agents
Sep 16, 2026Nilesh Verma, Nick Lim, Albert Bifet +1Tool-Using Agents ML Reproducibility
AI Institute, University of Waikato Hamilton 3216, New Zealand · LTCI, T´el´ecom Paris, Institut Polytechnique de Paris 19 place Marguerite Perey, 91120 Palaiseau, France
World Agent: Can Language Models Keep a World Running?
Sep 26, 2026Weixing Chen, Weipeng Zhang, Nan An +2World Model Evaluation World Models
Sun Yat-sen University · South China University of Technology · Guangdong Key Laboratory of Big Data Analysis and Processing +2

Jun 23, 2026J/K move · Enter open · S save

Qwen-AgentWorld: Language World Models for General Agents

Yuxin Zuo, Zikai Xiao, Li Sheng, Fei Huang+29

Qwen Team

LLM World Models Agentic RL World Model Learning World Model Evaluation

Open paper AskPDF arXiv

A world model predicts environment dynamics based on current observations and actions, serving as a core cognitive mechanism for reasoning and planning. In this work, we investigate how world modeling based on language models can further push the boundaries of general agents. (i) We first focus on building foundation models for agentic environment simulation. We introduce Qwen-AgentWorld-35B-A3B and Qwen-AgentWorld-397B-A17B, the first language world models capable of simulating agentic environments covering 7 domains via long chain-of-thought reasoning. Leveraging more than 10M environment interaction trajectories of 7 domains in real-world environments, we develop Qwen-AgentWorld through a three-stage training pipeline: CPT injects general-purpose world modeling capabilities from the state transition dynamics and augmented professional corpora, SFT activates next-state-prediction reasoning, and RL sharpens simulation fidelity through a tailored framework with hybrid rubric-and-rule rewards. To evaluate language world models, we present AgentWorldBench, a comprehensive benchmark constructed from real-world interactions of 5 frontier models on 9 established benchmarks. Empirical results demonstrate that Qwen-AgentWorld significantly outperforms existing frontier models. (ii) Beyond foundation models, we further investigate two complementary paradigms through which world modeling enhances general agents. First, as a decoupled environment simulator, Qwen-AgentWorld supports scalable and controllable simulation of thousands of real-world environments for agentic RL, yielding gains that surpass real-environment training alone. Second, as a unified agent foundation model, world-model training acts as a highly effective warm-up that improves downstream performance across 7 agentic benchmarks. Code: https://github.com/QwenLM/Qwen-AgentWorld

Abstract

Test-Time Scaling

Explore similar work

Qwen-AgentWorld: Language World Models for General Agents

TuiML: Machine Learning for AI Agents

World Agent: Can Language Models Keep a World Running?