Scaling Automatic Research Agents via World Models
Xiyuan Yang, Sheikh Sarwar, Jingru Cheng, Zhan Shi, Duanshun Li, Huiyuan Chen, Haiyang Zhang, Xing Fan, Chenlei Guo, Jingrui He, Zhenyu Liao
cs.LG
2026-08-13
WMRL replaces sandbox execution with a world model during GRPO, keeping a 10% real-execution anchor to debias and denoise. Training is about 3x cheaper; a post-trained 9B agent beats a 120B open-weight agent on held-out MLE-Dojo and DSBench.
RL for automatic-research agents splits into generation and execution. Generation batches. Execution needs an isolated sandbox and real machine time per candidate, so it saturates first as the trajectory count grows. Illinois and Amazon replace that execution with a few world-model forwards so the reward side can batch too.
The world model shares the agent backbone, stays frozen, and maps a task plus solution code to a predicted score. Standard GRPO draws n≥2 trajectories per task and subtracts the group mean. The world-model interface matches the real environment, so the update rule does not change; only the source of the scores does. Matching the backbone is deliberate: gains cannot be read as distilling a stronger teacher. Only the agent is updated. Using those scores alone injects bias b and noise ξ into the policy gradient. The convergence bound then carries extra O(B²) and O(σ²) terms, and the bias floor does not vanish with more steps.
WMRL keeps about 10% of groups as anchors and runs real execution on them. Online debiasing fits an isotonic map from world-model scores to true scores and refits as new pairs arrive. Inverse-variance denoising fuses the two gradient streams by estimated variance: noisier world-model updates get less weight; when residuals stay large because the outcome is not readable from the code, the update falls back toward ordinary GRPO. Both corrections tighten the bound, and the bias floor contracts with T.
Held-out MLE-Dojo and fully unseen DSBench, scored as real-leaderboard percentile, avg@8.
| Method | GPU-hours | MLE avg | DSBench avg |
| Qwen3.5-4B GRPO | 883 | 15.2 | 25.7 |
| 4B world model only | 269 | 12.9 | 23.1 |
| 4B WMRL | 286 | 16.4 | 28.8 |
| Qwen3.5-9B GRPO | 1174 | 18.8 | 31.2 |
| 9B WMRL | 349 | 21.6 | 32.8 |
| Kimi-48B-A3B | — | 8.1 | 17.3 |
| Nemotron-120B-A12B | — | 20.5 | 31.7 |
Against same-scale real-execution GRPO, compute drops about 3.1x and 3.4x while scores rise. The 4B agent beats the 48B open-weight agent and the 9B agent beats the 120B one; the latter margin is small and consistent. Pure world-model RL loses to real GRPO at both scales, so the corrections are load-bearing. Ablations show debiasing carrying more of the gain, with both pieces required for the full score.
Ablations on the 4B agent: neither correction 13.5 / 25.3 on MLE / DS; denoising only 14.9 / 26.2; debiasing only 15.7 / 28.1; both 16.4 / 28.8. The 9B row is 16.8/29.5, 18.0/31.2, 19.4/31.7, 21.6/32.8. The same recipe on MiniVLA-1B / LIBERO-Long moves overall success from 37.4 (SFT) and 38.3 (real GRPO) to 41.2. Robometer, an off-the-shelf VLM, scores eight frames per rollout as a dense progress reward; the environment’s sparse success at the end is the anchor. Sparse success or the world model alone barely moves the policy.
For these agents the expensive part is the environment, not the tokens. If outcomes are somewhat predictable from artifacts, and a thin ground-truth stream can be kept, a world model is a batchable reward machine. Research agents and VLA post-training both fit. Getting 4B/9B past much larger open-weight agents is an argument for this post-training over raw scale.
The world model is frozen and matched to the agent, so a new domain needs new calibration. Ten percent real execution remains. Image tasks stay weak (5.2 for 4B, 7.4 for 9B). The 4B text category is 0.5 below GRPO and 9B regression is 1.2 below. Train/test splits are author-drawn because the official MLE-Bench test split is flawed. The theory assumes smoothness and gradient domination. Leaderboard percentiles are avg@8; the main table does not report variance.