Predict the Reward, Not the World: World Reward Models for ML Agents

Yuxiao Lee; Zherui Yang; Fan Liu; Hao Liu

cs.AIOct 28, 2026

Predict the Reward, Not the World: World Reward Models for ML Agents

Authors: Yuxiao Lee, Zherui Yang, Fan Liu, Hao Liu

Organizations: The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China

Abstract

Search-based machine-learning (ML) agents improve a solution by writing candidate programs, executing them, and keeping the best. Execution gives the search reliable feedback, but a single run can take hours, and most runs end in programs the agent discards. Recent agents try to save these runs by letting a large language model act as a world model that predicts how each program would perform. However, a wrong prediction feeds the search false information. We argue that an agent does not need to predict the outcome of execution, only the decision that execution would support. We propose World Reward Models (WRMs), which learn from real execution results to choose directly which candidate will perform better, without predicting any intermediate outcome. We prove that a selector which predicts execution from the same inputs cannot outperform the best direct selector, that predicted outcomes add an error which can make a likelihood-trained world model prefer the worse program on every instance, and that cross-entropy on executed winners drives a WRM toward the optimum. Our WRM, ForeCritic, adapts a 9B typed decision model to 5k executed program pairs. On 12,438 held-out pairs from 26 ML tasks, it picks the better program 79.4% of the time, while general-purpose LLM judges reach 53--62%, and its probabilities are calibrated. Inside AIDE, ML-Master, and ForeAgent, it cuts mean task time by 4.4--7.5 times while keeping their mean scores within 0.01 or raising them.

Figures & tables

Figure 1: Three ways to choose which program to keep. (a) Executing every candidate gives real scores but costs hours per program. (b) A world model predicts execution instead; its predictions are right on about 60% of program pairs, and a wrong one misleads the search. (c) A World Reward Model learns from real executions to decide directly which program will score higher, and the agent runs only the program it keeps.

Figure 2: Predicting execution adds no information about the winner. (a) The predicted trace

z

is computed from

X

, so no world-model selector exceeds

\argmax\eta

( Theorem 3.1 ). (b) A world model’s decisions are those of a WRM with

w=\pi

; distilling it copies its accuracy, while training on executed outcomes targets the ceiling ( Propositions 3.3 and 3.4 ). (c) For value scores, predicting the run adds an error term that grows with the horizon

H

( Theorem 3.5 ).

Figure 3: A World Reward Model is a typed decision model trained on what the world returns. (a) Task, data report, and both programs form a shared state; a pointer head scores two short options against a decision token in one forward pass, with no generated token. (b) Program pairs labeled by their executed test scores train the WRM with cross-entropy. (c) Inside an agent, the WRM settles comparisons, dropping a loser only when confident; execution verifies the chosen program, and its score returns to the search tree.

Subset (pairs)	Kev-9B	ForeCritic	$\Delta$
Final test (12,438)	54.53	79.44	$+24.91$
no data report	54.46	78.55	$+24.09$
$\geq$ 1 unseen (5,996)	54.10	77.18	$+23.08$
2 unseen (501)	48.10	71.46	$+23.35$
Internal test (500)	52.60	78.20	$+25.60$
Task macro (26)	54.60	72.69	$+18.10$

Table 1: Training on execution outcomes lifts a typed selector by about 24 points, including on unseen programs. Pairwise accuracy (%) of the released Kev-9B and of ForeCritic, its trained version, under identical inputs and labels; “unseen” programs are absent from training.

\Delta

: paired difference, significant in every row under a paired bootstrap (over tasks for the macro row).

Figure 4: A WRM trained on real executions outperforms general-purpose judges. (a) Held-out program pairs. ForeCritic, Jev, and Kev-9B on the final test; Qwen3.5-9B on its 9,879-pair subset, where ForeCritic scores 79.9%. Hollow markers are reported by Zheng et al. (2026) on the full corpus, including how often the executed validation score picks the test-score winner. (b) 182 judgments of fresh programs that an agent wrote on ten MLE-Dojo tasks.

Figure 5: ForeCritic knows when it is right. Final test, 12,438 pairs. (a) Mean chosen probability against accuracy per bin; Kev-9B as one point (mean confidence, accuracy). (b) Accuracy of the decisions kept at threshold

\tau

against the fraction kept.

	Native score			Task time (s) $\downarrow$
Agent	Orig.	Jev	ForeCritic	Orig.	Jev	ForeCritic
AIDE	0.600	0.499	0.599	3,164	738	725
ML-Master	0.616	0.484	0.609	3,711	839	497
ForeAgent	0.527	0.472	0.607	4,111	2,094	884

Table 2: An accurate WRM keeps or raises agent scores at a fraction of the time; a near-chance one lowers them. Mean native score and mean task time over ten MLE-Dojo tasks with each agent’s original selector, with Jev, and with ForeCritic, under the same agent configuration; one scored run per task, timed from start to submission. Native metrics differ in scale and direction (7 of 10 are higher-is-better; Figure 11 ).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: In the idealized Gaussian model, breadth compensates for accuracy only slowly. (a) Lemma B.2 maps pairwise accuracy

p

to correlation

\rho

. (b) Expected gain

\rho\,e_{K}

in units of

\sigma

( Theorem B.3 ) for illustrative

p=0.615

and

p=0.794

: the

p=0.615

judge needs

K=907

candidates to match the

p=0.794

judge at

K=8

. Curves are analytic and are not fitted to the agents studied here.

Figure 7: The expected maximum of

K

standard normals grows like

\sqrt{\ln K}

e_{K}=\int zK\phi(z)\Phi(z)^{K-1}\,dz

by numerical integration, with

\phi

and

\Phi

the standard normal density and distribution function, against the bound

\sqrt{2\ln K}

of Theorem B.3 . Multiplying

K

by 128, from 8 to 1024, multiplies

e_{K}

by 2.3. The marked points give the catch-up example:

e_{907}

\rho_{R}/\rho_{W}=2.26

times

e_{8}

Figure 8: On the 9,879 pairs common to all evaluations, only the execution-adapted models clearly exceed chance. Pairwise accuracy on the same pairs and labels; the models differ in weights, inputs, and training. Kev fine-tune v1 is an earlier adaptation with a different recipe and split. Labels on the right give ForeCritic’s paired advantage, which is significant against every model.

Figure 9: Temperature scaling refines an already calibrated ForeCritic and repairs Kev-9B’s probabilities but not its accuracy; ForeCritic’s confidence trades coverage for accuracy. Final test, 12,438 pairs. (a) NLL, Brier score summed over the two classes, and ECE over ten bins, before and after temperature scaling; accuracy is unchanged by scaling (54.53% for Kev-9B, 79.44% for ForeCritic). (b) Fraction of decisions kept and their accuracy when only decisions with chosen-class probability at least

\tau

are accepted; thresholds were not tuned on these labels.

Figure 10: On fresh agent candidates, ForeCritic is more accurate in each presentation order but less order-consistent than DeepSeek. 91 pairs from ten MLE-Dojo tasks. (a) Accuracy in each order and, for ForeCritic, after averaging the two orders’ probabilities. (b) Each pair judged in both orders is right twice, right once (the two orders disagree), or wrong twice.

Figure 11: With ForeCritic as selector, the three agents are within 4% of each other on five of ten tasks, and a few long tasks dominate task time. One scored run per task. (a) Relative gap between each agent’s native score and the best of the three on that task, taking the metric direction into account (0 is the best); under each task are its metric, direction, and best score, and the bottom row gives the mean native scores of Table 2 . (b) Task time in minutes.

Preference Learning

Preference learning infers people’s values and trade-offs from choices, comparisons, rankings, and feedback to guide predictions and decisions.

See all 152 papers

MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model
Oct 27, 2026Brian Ko, Han Wang, Hejia Geng +2Test-Time Scaling World Model Evaluation
University of Washington · University of Illinois at Urbana-Champaign · University of Oxford +2
Language Models as AI Research World Models
Oct 8, 2026Zijun Wang, Zewen Liu, Minhua Lin +10LLM World Models Autonomous Research Agents
Amazon · UC Santa Cruz · Emory University +2
Agent-Editing World Model: Rethinking World Modeling for LLM Agents
Sep 23, 2026Shuang Sun, Guoxin Chen, Fanzhe Meng +6LLM World Models World Models
Gaoling School of Artificial Intelligence, Renmin University of China

Oct 27, 2026J/K move · Enter open · S save

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

Brian Ko, Han Wang, Hejia Geng, Li Erran Li+1

University of Washington · University of Illinois at Urbana-Champaign · University of Oxford+2

Test-Time Scaling World Model Evaluation

Open paper AskPDF arXiv

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/.

Abstract

Preference Learning

Explore similar work

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

Language Models as AI Research World Models

Agent-Editing World Model: Rethinking World Modeling for LLM Agents