PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuxiao Lee, Zherui Yang, Fan Liu, Hao Liu

arXiv:2610.38439v1cs.AIcs.CL

TL;DR

Search-based ML agents rely on costly executions for reliable feedback, while predicting execution can mislead selection. The paper introduces World Reward Models, which learn to choose directly from real execution results. ForeCritic performs better than general-purpose LLM judges on held-out program pairs and speeds up three agents while maintaining or improving their scores.

  • Problem

    Search-based ML agents use reliable but costly execution to select programs, while world-model predictions can mislead search when wrong.

    Source: §1 · p. 1

  • Method

    World Reward Models learn from real execution outcomes to choose directly among candidate programs without predicting intermediate execution outcomes.

    Source: §1 · p. 2

  • Results

    79.4%: ForeCritic picks the better program on held-out pairs, versus 53–62% for general-purpose LLM judges; mean task time falls 4.4–7.5× with similar or better agent scores.

    Source: §1 · p. 2, §8 · p. 8

  • Takeaways & Limitations

    ForeCritic can replace most executions in AIDE and ML-Master and the LLM judge in ForeAgent while preserving or improving mean scores.

    Source: §1 · p. 2

    The selector learns from executions on a family of ML tasks and settles one comparison at a time; transfer to new tasks and domains and ranking many candidates in one pass remain part of the broader agenda.

    Source: §7 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

Search-based ML agents gain reliable feedback from execution, but runs are costly and world-model prediction errors can mislead search. The paper proposes World Reward Models that learn to choose directly from real execution results, and evaluates ForeCritic in agents and on held-out program pairs.

1. Introduction

  • A single ML-bench run can take hours, and most runs end in programs the agent discards, despite execution providing reliable scores from data.AIDE and ML-Master execute candidate programs and use validation scores to decide what to refine next.Source: §1 · p. 1
  • With reasoning enabled, DeepSeek-V3.2 and GPT-5.1 pick the better program on 61.5% and 58.8% of pairs, and errors can promote worse programs or prune better ones.A fast selector that is right on 56% of pairs lowers all three tested agents’ final scores.Source: §1 · p. 2
  • A World Reward Model reads the task and candidate programs to estimate which would win, learning from real runs without predicting any part of execution.During search, it takes the place of most executions.Source: §1 · p. 2
  • A selector that predicts execution from the same inputs cannot beat the best direct selector; prediction adds error, while cross-entropy on executed winners bounds a WRM’s accuracy gap.The paper also shows that likelihood-trained execution logs can make a world model prefer the worse program on every instance.Source: §1 · p. 2
  • Mean task time drops 4.4–7.5× with ForeCritic, while ForeAgent’s mean score rises from 0.527 to 0.607 and AIDE and ML-Master stay within 0.01 of their original scores.On 12,438 held-out pairs from 26 tasks, ForeCritic picks the better program 79.4% of the time, versus 53–62% for general-purpose LLM judges.Source: §1 · p. 2

2. Choosing Before Executing

This section frames pre-execution selection as choosing among candidate programs using the task context and the proposal policy. It distinguishes direct selectors from selectors that predict execution, and defines World Reward Models as direct selectors.

2. Choosing Before Executing

  • A decision instance X consists of a context and K candidates; in ML engineering, the context includes the task description and data report, and each candidate is a complete program.Executing candidates produces outcomes and rewards, with Y denoting the highest-reward candidate and ties broken uniformly at random.Source: §2 · p. 2
  • Accuracy and value can disagree: a candidate that wins more often need not have the higher expected reward.The section’s example contrasts a candidate with constant reward 0.6 against one that scores 0.61 with probability 0.51 and otherwise scores zero.Source: §2 · p. 2
  • A world-model selector first predicts execution outcomes and then applies a decision rule to those predictions.The definition covers simulate-then-evaluate planning and LLM reasoning about programs; both consult only X.Source: §2 · pp. 2–3
  • The no-new-evidence assumption makes the predicted trace and decision randomness independent of real outcomes and winner tie-breaking, given X.The assumption fails when the procedure gathers new evidence, such as by training on a data subsample or querying the environment; those steps count as executions.Source: §2 · p. 3
  • A preference WRM maps X to candidate-win probabilities, while a value WRM maps X to candidate scores; either selects the argmax without generating an intermediate outcome.A preference WRM is calibrated when its stated probabilities match conditional winner frequencies.Source: §2 · p. 3

3. What Predicting Execution Can and Cannot Add

The section argues that predicting execution cannot add information to the decision supported by the same inputs, while errors in predicted runs can hurt selection. It contrasts distilling a world model with training directly on executed winners, then analyzes value-error growth and a case where likelihood training selects the worse candidate.

3.1. A prediction carries no new evidence

  • A predicted trace cannot beat the Bayes-optimal WRM on the same inputs, and more rollouts, reasoning tokens, or parameters cannot cross that ceiling.The theorem compares selectors with the same information, not networks of equal size; which route gets closer to the ceiling is an empirical question.Source: §3.1 · p. 3

3.2. Same decisions, different targets

  • World-model selectors and WRMs can express the same decisions, but distillation copies the selector’s errors while executed-winner cross-entropy targets η.At the population level, the accuracy gap is at most the square root of twice the excess cross-entropy; the bound makes no claim about sample size.Source: §3.2 · pp. 3–4

3.3. Where predicting the run costs accuracy

  • Predicted value error is bounded by min{1, Hεsim} + εeval, while a direct value WRM with sup-error εw has regret at most 2εw.The simulation term grows with the execution horizon H; when the top-two value gap exceeds 2ε, choosing by predicted values loses nothing.Source: §3.3 · p. 4
  • Likelihood training can make a world model select the lower-value candidate on every instance when decision-irrelevant log bits outweigh reward differences.For q = (0.8, 0.3), m = (0.1, 0.9), and d = 10, the model incurs regret 0.5; in an idealized Gaussian model, a 61.5%-accurate judge needs 907 candidates to match a 79.4%-accurate judge seeing 8.Source: §3.3 · p. 4

4. Building and Using a World Reward Model

The paper builds a pairwise World Reward Model that learns from executed program comparisons, then integrates it into search-based ML agents. It describes the model interface, training data and objective, and how the WRM replaces candidate executions or an LLM judge during search.

4. Building and Using a World Reward Model

  • The proposed pairwise WRM learns to compare ML programs from executed outcomes and is plugged into search-based agents.Its training objective is cross-entropy on executed winners, and the model settles comparisons inside an agent’s search.Source: §4 · p. 4, Figure 3 · p. 5

4.1. A WRM as a typed decision model

  • A typed decision model reads the task, data report, and both programs, then scores two options in one forward pass without generating tokens.A pointer head uses the decision-token state and each option’s closing-token state; the fitted temperature calibrates the probabilities.Source: §4.1 · p. 5

4.2. Training on what the world returns

  • ForeCritic trains on 5,000 executed program pairs drawn from a corpus of 18,438 pairs across 26 MLE-bench tasks.The 12,438-pair final test is held out; the WRM minimizes cross-entropy on executed winners and fits temperature on a separate calibration partition.Source: §4.2 · p. 5

4.3. Plugging a WRM into search-based agents

  • In AIDE and ML-Master, the WRM replaces candidate execution with win probabilities; in ForeAgent, it replaces the LLM judge and drops a loser only above τ = 0.7.The agents still execute the submitted program, with repairs if needed; AIDE uses log-odds in node scores, while ML-Master backs up win probabilities as rewards.Source: §4.3 · p. 6

5. Experiments

Across the held-out experiments, execution-trained ForeCritic improves pairwise selection and supports faster ML-agent runs with maintained or higher scores. The evaluation also tests generalization, calibration, and comparisons with general-purpose judges.

5.2. Q1: Learning from execution makes the selector accurate

  • Training on execution outcomes raises final-test pairwise accuracy from 54.53% to 79.44%, a paired gain of 24.91 points.The gain is 18.10 points with every task weighted equally and remains significant under whole-task resampling; on 501 pairs with neither program seen in training, ForeCritic reaches 71.46% with a 23.35-point gain. Removing the data report costs 0.89 points while shortening the input by 63.9%.Source: Table 1 · p. 6, §5.2 · pp. 6–7

5.1. Setup

  • The final test holds out 12,438 program pairs; significance uses paired bootstrap at 5%, with task resampling accounting for pairs that share a program.The main comparison uses ForeCritic and its released Kev-9B initialization under the same inputs, labels, and architecture.Source: §5.1 · p. 6

5.3. Q2: Comparison with general-purpose LLMs

  • ForeCritic is right on 79.9% of the 9,879 pairs shared with Qwen3.5-9B, versus 53.1% for Qwen3.5-9B with reasoning and 61.5% for DeepSeek-V3.2.On 91 pairs of ForeAgent-written programs, ForeCritic is right on 70.3% of judgments versus 63.7% for DeepSeek-V4.1-Flash, a difference the small sample does not make significant.Source: Figure 4 · p. 7, §5.3 · p. 7

5.4. Q3: The probabilities are calibrated and informative

  • ForeCritic’s unscaled ECE is 0.025; thresholding at 0.9 retains 38.3% of pairs at 95.2% accuracy.Temperature scaling lowers NLL from 0.4364 to 0.4332 and ECE to 0.016 without changing decisions.Source: Figure 5 · p. 7, §5.4 · p. 7

5.5. Q4: The WRM inside three ML agents

  • Task time falls 4.4× in AIDE and 7.5× in ML-Master, with mean scores within 0.01; ForeAgent is 4.7× faster and its mean score rises from 0.527 to 0.607.ForeCritic comparisons take under 3% of task time; each agent comparison averages one run per task.Source: Table 2 · p. 7, §5.5 · p. 8

6. Related Work

The paper situates World Reward Models among work on learned world models, value learning, reward models, verifiers, and performance prediction. Its distinction is direct comparison of free-form programs for one task within agents that iterate using execution feedback.

Related work: World models for planning and LLM agents

Related work: Learning values instead of models

  • Prior value-learning work learns action values without a model, while value-aware modeling and value equivalence focus model accuracy where values depend on it [Grimm et al., 2020].The paper extends these ideas to one-step program selection and notes that likelihood-trained models can hurt control [Lambert et al., 2020].Source: §6 · p. 8

Related work: Reward models, verifiers, and judges

Related work: Predicting program performance

7. Limitations

The paper’s current WRM is limited to execution-trained selection among programs on a family of ML tasks, leaving broader capabilities and evaluations for future work.

Limitations

  • The current selector learns from executions on ML tasks and compares one pair at a time; transfer across domains, ranking many candidates, and continual learning remain future directions.The authors also leave matched-scale comparisons of training signals and interfaces, finite-model and finite-data theory, and fixed-budget final-reward evaluation across agents for future work.Source: §7 · p. 8

8. Conclusion

The conclusion presents World Reward Models as a way to replace costly execution-based selection with direct, execution-trained comparisons, while retaining execution to verify chosen programs.

Conclusion

  • 79.4%: ForeCritic picks the better of two ML programs, well above general-purpose LLM judges, and makes three agents 4.4–7.5× faster at similar or better scores.It learns the probability that one program beats another from real executions, while execution remains the verification step.Source: §8 · p. 8

Impact Statement

The impact discussion recommends retaining execution as verification and routing uncertain choices to execution or a human. The proof appendices establish a total-variation fact, characterize fixed-information selection, relate world-model selectors to WRMs, and connect cross-entropy to selection accuracy.

Impact Statement

  • A confidently wrong selector may steer an agent toward harmful actions when actions have consequences outside a sandbox; the authors recommend routing uncertain choices to execution or a human.WRMs can lower compute and energy spent executing poor candidates in ML-engineering agents.Source: Impact Statement · p. 8

A. Proofs

  • For probability measures, expectations of any measurable function valued in [0, 1] differ by at most their total-variation distance.The appendix states this inequality as a proof tool.Source: A. Proofs · p. 11

A.1. Fixed-information ceiling

  • The fixed-information proof characterizes zero selection gap: a selector must choose only maximizers of the relevant value or accuracy function almost everywhere.An argmax rule based on the inputs attains the value or accuracy optimum without observing outcomes.Source: §A.1 · p. 12

A.2. Behavioral equivalence and distillation

  • A sampled WRM can match a world-model selector’s decision distribution, yielding equal value and accuracy; conversely, a WRM can be represented as a world model.Distillation uniquely matches the teacher’s decision distribution, but inherits its errors; taking argmax can change accuracy.Source: §A.2 · p. 12
  • Distilling a teacher’s distribution does not guarantee matching its accuracy when replacing sampling with argmax, and it inherits teacher errors.The proof gives two-candidate examples where argmax is less accurate in one case and more accurate in another.Source: §A.2 · p. 12

A.3. Cross-entropy controls selection accuracy

  • The cross-entropy proof treats excess loss as expected KL divergence and uses Jensen’s inequality to average instance-wise accuracy gaps.Its intermediate steps compare the predicted and optimal choices on each instance.Source: §A.3 · p. 13

A.4. Error decomposition

The error analysis bounds how simulation and evaluation inaccuracies affect selection, then shows how likelihood fitting can favor a worse program. Under an idealized Gaussian model, improving judge accuracy can be far more effective than screening many more candidates.

A.4. Error decomposition

  • Simulation error accumulates with horizon, giving total-variation error at most min(1, Hεsim); adding evaluator error bounds value error by min(1, Hεsim) + εeval.The proof compares the real and modeled outcome distributions before bounding the resulting value difference.Source: §A.4 · pp. 13–14
  • A direct value WRM with uniform value error εw has regret at most 2εw.The bound applies when its estimated values differ from true values by at most εw for every instance and candidate.Source: §A.4 · p. 14

A.5. Objective mismatch in a restricted family

  • Likelihood fitting can select the worse program on every instance when d(m2 − m1) > q1 − q2, despite an exact evaluator.In the stated construction, reward-only fitting instead recovers estimates (q1, q2) and selects a1.Source: §A.5 · p. 14
  • For q = (0.8, 0.3), m = (0.1, 0.9), and d = 10, likelihood fitting selects a2 with regret 0.5 and 0.25 winner accuracy, versus 0.75 for the Bayes preference selector.The construction concerns population objectives in the tied-rate family.Source: §A.5 · p. 14

B. Idealized Gaussian Screening

  • Under the Gaussian screening assumption, pairwise accuracy determines correlation through p = 1/2 + (1/π) arcsin ρ, and expected selected-value gain scales with ρeK.Here eK is the expected maximum of K independent standard normals.Source: B. Idealized Gaussian Screening · p. 15
  • A judge with pairwise accuracy 0.615 needs 907 candidates to match a 0.794-accuracy judge screening 8 in the idealized model.The curves are analytic and are not fitted to the agents studied.Source: Figure 6 · p. 15, B. Idealized Gaussian Screening · p. 15

B.1. Proofs

The proofs connect pairwise accuracy to correlation and derive the expected gain from screening under Gaussian assumptions. They also establish a knockout-selection result and give numerical examples of how slowly breadth compensates for lower accuracy.

B.1. Proofs

  • The proof relates pairwise accuracy to correlation by applying Sheppard’s formula to the value and score differences, yielding p = 1/2 + (1/π) arcsin ρ.The differences form a centered bivariate normal vector with correlation ρ.Source: §B.1 · pp. 15–16
  • The expected selected value is μV + ρσeK, and eK is nondecreasing with candidate count.For two candidates, e2 = 1/√π; the numerical section also compares computed values with the bound √(2 ln K).Source: §B.1 · p. 16, §B.2 · p. 17
  • For K = 2^m candidates, if the best candidate wins each surviving-round match with conditional probability p independently, its knockout selection probability is p^m.This conditional match accuracy generally differs from average pairwise accuracy over all pairs.Source: §B.1 · pp. 16–17

B.2. Numerical values

  • Under the idealized Gaussian model, a judge with pairwise accuracy 0.615 needs at least 175 candidates, and exactly 907 numerically, to match a 0.794-accuracy judge screening 8.The numerical check verifies the formulas, not their fit to real agents.Source: §B.2 · p. 17

C. Data, Training, and Inference Details

The paper describes ForeCritic’s data split, training, and calibration, then reports comparisons, probability quality, transfer checks, and agent integrations. The supplied agent results include mean native scores and task-time examples, while noting that each configuration has one scored run per task.

C. Data, Training, and Inference Details

  • The corpus was split into 5,000 training pairs and separate development, calibration, internal-test, and final-test partitions; no source pair or identical unordered program pair crosses partitions.Tasks are shared across partitions, and 5,996 final-test pairs contain at least one program absent from training.Source: C. Data, Training, and Inference Details · p. 17

D.1. Shared-pair comparison with Qwen and an earlier adaptation

  • On 9,879 pairs shared across evaluations, the Section 4.2 recipe reaches 79.89% accuracy, versus 60.82% for the earlier adaptation.Reasoning raises Qwen3.5-9B from 52.82% to 53.12%; the Figure 8 comparison covers the same pairs and labels.Source: Figure 8 · p. 18, §D.1 · p. 18

D.2. Probability quality

  • ForeCritic’s accuracy stays 79.44% after temperature scaling, while its NLL changes from 0.4364 to 0.4332 and ECE from 0.0251 to 0.0163.Raising the acceptance threshold from 0.5 to 0.95 reduces accepted decisions from 100% to 25.2% and raises accuracy from 79.4% to 97.3%.Source: §D.2 · p. 18, Figure 9 · p. 19

D.3. Transfer details

  • On transfer candidates, a 0.7 confidence threshold keeps 90 of 182 judgments, with 72 correct (80.0%); ForeCritic is consistent across orders on 69 of 91 pairs.Averaging the two order probabilities yields 63 of 91 correct (69.23%), an insignificant change of −1.10 points from mean single-order accuracy.Source: §D.3 · p. 19, Figure 10 · p. 19

E. Agent Integration Details

  • Across 30 ForeCritic submissions, mean native scores are 0.5994 for AIDE, 0.6088 for ML-Master, and 0.6068 for ForeAgent.A few long tasks account for much of task time; each configuration has one scored run per task, and the mean-score differences are not tested for significance.Source: E. Agent Integration Details · p. 20, Figure 11 · p. 21

Related papers

Loading 2610.38439v1…