Robotics · 2026-10-30
Predict the Reward, Not the World: World Reward Models for ML Agents
Yuxiao Lee, Zherui Yang, Fan Liu, Hao Liu
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors propose World Reward Models (WRMs) that learn to choose which candidate ML program will score higher using only real execution outcomes, without predicting intermediate execution traces. They prove limits of predicting execution and show how cross-entropy on executed winners can drive a WRM toward optimal selection. Their ForeCritic (from Kev-9B) improves pairwise selection accuracy and speeds up search-based agents while keeping scores similar.
Why it matters
In search-based ML agents, executing every candidate is slow and many runs are discarded. WRMs aim to replace execution with a direct decision model trained on executed winners, avoiding failure modes of world models that predict execution and propagate errors into pruning. The paper also provides theoretical arguments and empirical results showing time reductions with controlled score changes.
Method
- Define a World Reward Model (WRM) as a selector that maps (task, candidate programs) directly to win probabilities for each program, without generating intermediate execution outcomes.
- Train a pairwise preference WRM on program pairs labeled by which executed program achieved the better test score, using cross-entropy on the executed winner.
- Plug the WRM into search-based ML agents by using it in place of execution scoring for candidate selection (including confidence-gated elimination).
Abstract (from arXiv)
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.
Related papers
- MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model
- ANTE: Anticipate, Then Execute -- Learning Internal World Models for Coding Agents from Verified Beliefs
- How Much Can Best-of-N Help a World Action Model? Measuring Selectable Headroom and Realizing It in Closed Loop
- Tev: Scaling Long-Horizon Agentic Environments with Dense Progress Rewards
- Self-Improving Agents with World Model Predictive Paths
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First announced by arXiv: 2026-10-30 · Indexed by Pub-AI: 2026-10-30 UTC
Predict the Reward, Not the World: World Reward Models for ML Agents
First announced by arXiv: 2026-10-30 · Original record ↗
- arxiv
- 2610.38439
Explicit source versions (1)
- v1 · Date not supplied
- Preprint available
Source date: 2026-10-30 · First detected: 2026-10-30 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · RSS guid=oai:arXiv.org:2610.38439v1; arxiv:announce_type=new; item pubDate (official metadata)
Abstract · arxiv
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.