Scaling Automatic Research Agents via World Models

Scaling Automatic Research Agents via World Models. By Xiyuan Yang and ten others, from the University of Illinois Urbana-Champaign and Amazon. Published on the twenty-ninth of August. Its code repository is worked out: about sixty-five source files, with a readme and a licence.

Automating empirical research is an old ambition of AI, and automatic research agents, AutoResearch agents, are bringing it within reach. Modern LLMs can implement solutions on their own and learn from what happens when those solutions are executed. Behind those gains, post-training plays a central role, and reinforcement learning above all.

Scaling that reinforcement learning, the authors argue, runs into a fundamental tension. Every AutoResearch trajectory has two components: the agent generating, and the environment executing what it generated. The two scale in very different ways. Generation shares compute, because many requests can be batched together on the same hardware. Execution cannot be shared like that. Each run occupies its own exclusive sandbox and consumes real machine time. So as trajectories grow longer, environment execution comes to dominate the cost of training, and it becomes the bottleneck that limits how far the method can scale.

World Model RL, WMRL for short, removes that bottleneck by replacing environment execution with a world model. The expensive, unshareable half of each trajectory is handed to a learned model instead of being run for real. Generation continues exactly as before. What changes is where the feedback comes from.

The authors are direct about the weakness this introduces. A world model can be imperfect, and its rewards are corrupted in two distinct ways, by bias and by noise. Bias is a consistent lean in one direction; noise is scatter around the truth. WMRL is equipped with one mitigation for each. Online Debiasing offsets the bias. Inverse-Variance Denoising suppresses the noise. Rather than treating the world model's reward as correct, the method corrects it along both axes before learning from it.

On the theoretical side, the authors prove that both mitigations strictly improve the convergence guarantee.

Empirically, WMRL speeds training up by three to four times across various tasks and agent scales, while also exceeding standard reinforcement-learning baselines on performance. The post-trained agents, at four billion and nine billion parameters, outperform much larger open-weight agents of forty-eight billion and one hundred and twenty billion parameters on held-out benchmarks.

The method also reaches beyond research agents. WMRL transfers to post-training embodied VLA policies, which the authors present as a demonstration that it generalizes.

The diagnosis drives everything. In agent training, generating text is cheap to share and running experiments is not, so the expensive half gets replaced by a learned stand-in. What makes the substitution usable is taking seriously that the stand-in is imperfect, and correcting its bias and its noise separately rather than trusting it outright. The reward is then cheap enough to scale and corrected enough to learn from.

Scaling Automatic Research Agents via World Models. By Xiyuan Yang and ten others, from the University of Illinois Urbana-Champaign and Amazon. Published on the twenty-ninth of August. This is the long version: the whole paper rather than the abstract.

Reinforcement learning is a central way to post-train research agents: language models that take a machine learning problem, write a solution, run it, read the result and try again. That demands a vast number of trajectories, and the authors point to a tension inside every one of them. A trajectory has two parts, and they scale differently. Generating the agent's code is served by inference engines that batch many requests together, so an extra trajectory costs almost nothing. Executing that code cannot be batched. Every candidate solution needs its own isolated sandbox that trains a model on a real GPU, so execution cost grows linearly, and it is execution that hits capacity first.

The training recipe underneath is group-relative policy optimization. For each task the agent produces a group of eight attempts, each is scored, and each score is compared against the group's average. The score here is the solution's percentile on the real Kaggle leaderboard for that competition. The paper's proposal is to replace execution with a world model: a language model that reads the task and the agent's code and predicts what the sandbox would have reported. That takes a few forward passes, so it batches like generation. The price is error, and the paper splits that error in two. The bias is the systematic part that survives averaging. The noise is the zero-mean remainder. In the convergence proof, noise inflates the variance of each update, damped by the step size. Bias is worse: it leaves a floor on the final gap to the optimum that no amount of training removes.

The world model is deliberately modest. It is the same backbone as the agent, never fine-tuned, prompted only, and frozen throughout training. That rules out quietly distilling a stronger model. Its prompt splits the job into eight parallel single-aspect checks: from imports and API calls against the installed library versions, through file and column references and the compute budget, to submission validity and solution quality. It is told to default to pass, because a false failure on working code is worse than a missed error, and its failure messages imitate the real sandbox's traceback so the agent can fix its code. Failures cap the score, and a valid solution is scored by the quality check.

The corrections need ground truth, so the method keeps a thin stream of it. About ten percent of groups, at least one per step, are also executed for real. These anchor groups give pairs of predicted and true scores. Online Debiasing fits a monotone map from one to the other by isotonic regression: first fit after two hundred pairs, then refit as new ones arrive to follow the drift. Inverse-Variance Denoising then merges the two gradient streams, one scarce and clean, one abundant and noisy, weighting each by the inverse of its variance, which gives lower variance than either alone. The anchor weight rises with the measured disagreement between calibrated and true scores, clipped between one and four. When a task's outcome hinges on randomness no reading of the code can reveal, the disagreement stays large, the anchors carry the gradient, and the method falls back towards ordinary training. In the theory, the bias floor now shrinks as training proceeds and vanishes in the limit.

The agents are Qwen three point five models at four and nine billion parameters. They train on forty-five small Kaggle competitions and are tested on fourteen held-out ones, each larger than every training task in its category. A second suite, DSBench, adds sixty more. Against real-execution training, the nine billion agent's compute falls from one thousand one hundred and seventy-four GPU-hours by three point four times, and its held-out average still rises, from eighteen point eight to twenty-one point six percentile. Both post-trained agents beat far larger open-weight models on both averages.

Pure world-model training is cheaper still, but its averages land below real execution at both scales and on both suites. The comparison that tests the thesis most directly is the ablation, where every run sees the same two reward streams in the same ratio. Feeding both to training with no correction also stays below real execution on every column. Denoising alone adds under two points. Debiasing alone adds more, up to about three, which is what the theory predicts, since bias enters the bound at full size while noise is damped by the step size. Together they add between two point nine and four point eight, more than either alone. The recipe also transfers to post-training a robot manipulation policy.

The paper has no limitations section. What its evidence supports is narrower than its framing. Each score is averaged over eight attempts at evaluation, but no repeated training runs or error bars are reported, and the margins over real-execution training are modest: gains of up to three point one points, with two category columns where the new method scores lower. The held-out MLE-Dojo set is fourteen competitions. The theory also assumes the bias is a monotone distortion of the true score, which is what makes a monotone fit able to learn it. And the speedup depends on the world model being predictive at all: on tasks where it is not, the design falls back towards full-cost training by construction.

The lasting idea is about where the cost of agent training really sits. When generation is cheap and execution is not, a learned stand-in for the environment can remove the bottleneck, but only if something keeps it honest. Here that is a thin stream of real results, used to recalibrate the stand-in and to decide, step by step, how much to trust it.