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
World Model RL (WMRL) accelerates AutoResearch agent training by 3–4× by replacing expensive environment execution with a learned simulator.
How can we scale AutoResearch agents by replacing expensive real-world environment execution with a world model while correcting for the model's inherent bias and noise?
Training autonomous research agents requires massive environment execution, but each experiment must run in an isolated sandbox, creating a linear scaling bottleneck that cannot be amortized like model generation. World Model RL (WMRL) replaces these expensive environment calls with a learned world model that simulates outcomes, while using a small, concurrent stream of real execution to anchor and correct the model's inherent bias and noise. This approach accelerates training by 3–4× compared to standard reinforcement learning, while consistently outperforming full-environment baselines on held-out benchmarks.
Paper Primer
WMRL treats the world model as a surrogate environment: it predicts the execution outcome of an agent's solution, allowing the training pipeline to scale compute-efficiently. To handle the inevitable inaccuracies of this simulation, the method uses an anchor signal—a small fraction of groups graded by both the world model and real execution—to perform online error correction.
The core mechanism is a dual-correction strategy: Online Debiasing uses isotonic regression to map biased world model scores toward the ground truth, while Inverse-Variance Denoising fuses the world model and anchor gradients based on their relative reliability. This is like a navigator using a high-frequency, noisy sensor alongside a low-frequency, precise GPS to maintain a stable trajectory.
WMRL significantly reduces training compute while improving final agent performance.
On AutoResearch tasks, WMRL achieved a 3.1–3.4× speedup in training compute compared to standard reinforcement learning (GRPO) while exceeding its performance by up to 3.1 points.
Post-trained agents outperform significantly larger, off-the-shelf models.
A 4B parameter agent trained with WMRL surpassed a 48B parameter agent, and a 9B agent surpassed a 120B agent on held-out benchmarks.
Why is environment execution a bottleneck for these agents, but not for standard LLM training?
In LLM generation, batching allows multiple trajectories to share compute. In AutoResearch, each experiment requires an isolated sandbox (e.g., a Docker container) to run code, meaning execution cost scales linearly with the number of trajectories rather than being amortized.
Does this method require a perfect world model to function?
No. The authors prove that the world model's bias and noise introduce error terms into the convergence bound, which the Online Debiasing and Inverse-Variance Denoising mechanisms provably reduce, allowing the agent to converge to the same optimum as training with real execution.
WMRL demonstrates that RL for agentic tasks can be decoupled from the cost of environment execution, provided a small, consistent stream of ground truth is available to anchor the surrogate model.
The AutoResearch Bottleneck
AutoResearch training stalls due to costly environment execution, prompting a shift to world‑model simulation.
Training AutoResearch agents requires generating many candidate solutions and executing each in a real environment. While generation can be batched and thus cheap, execution needs an isolated sandbox per solution, making environment execution dominate the training cost as trajectories increase.
The training pipeline stalls because each generated solution must be run in a real sandbox, and that cost grows linearly with the number of trajectories.
**Figure 1.** AutoResearch trajectories scale asymmetrically, making environment execution the bottleneck, and our WMRL removes it. (a) Per research question, the agent proposes a group of solutions $S_i$, graded by execution into rewards $r_i$, and RL demands a massive number of such trajectories. (b) Per trajectory, generation amortizes compute via batching, but execution needs one isolated sandbox per solution. (c) Execution compute thus hits capacity first in traditional RL, whereas WMRL scales without this limit.
Two research questions arise: (i) can we replace the expensive environment with a fast, scalable signal, and (ii) what cost does this surrogate incur? The answer to (i) is World Model RL, which simulates execution outcomes; (ii) is addressed by Online Debiasing and Inverse‑Variance Denoising, which correct bias and noise introduced by the world model.
Shifting from real‑world execution to world‑model simulation removes the scaling bottleneck.
The WMRL Framework
Introduce WMRL, correct world‑model bias and variance with a thin stream of real rewards.
Standard RL with group‑relative policy optimization (GRPO) requires executing each candidate solution in a real environment, which dominates training cost.
Replace the expensive environment execution with a learned world model that predicts the reward, while keeping a tiny “anchor” stream of real executions to correct systematic errors.
Compute raw advantages using real scores for the anchor: $A_1 = 0.8 - 0.65 = 0.15$.
Compute surrogate advantages for all groups using $\hat{r}$: $\hat{A} = [0.75-0.70, 0.55-0.70, 0.85-0.70, 0.65-0.70] = [0.05, -0.15, 0.15, -0.05]$.
Form the gradient estimate $\hat{g}$ as the average of $A_i s_i$ (here we treat $s_i$ as unit vectors for illustration).
Because the anchor provides the true $A_1$, the bias $b$ can be estimated as $\hat{r}_1 - r_1 = -0.05$, which will be used later for correction.
The anchor group supplies a concrete error signal that lets the algorithm learn to shift all $\hat{r}$ values toward the true scale, despite most groups never seeing a real rollout.
Using the world model introduces two new error sources: a systematic bias and stochastic noise, which we now address.
Fit a monotone mapping $\hat{f}$ on the paired scores from anchor groups so that the world‑model predictions are calibrated back to the true reward scale.
Apply $\hat{f}$ to a world‑model score $\hat{r}=0.85$ → calibrated $\hat{f}(0.85)=0.90$.
Form the calibrated advantage $\hat{A}=0.90 - \frac{1}{4}(0.90+0.60+0.80+0.70)=0.90-0.75=0.15$.
Use this calibrated advantage in the gradient estimate, reducing the systematic bias.
Even a tiny linear correction can eliminate the constant offset that would otherwise accumulate over many updates.
Combine the scarce, low‑variance gradient from real executions with the abundant, higher‑variance gradient from the world model, weighting each by the inverse of its variance.
Compute $\rho = V_{W_M}/V_E = 0.0016/0.0004 = 4$.
Fuse: $\hat{g} = \frac{4\,g_E + g_{W_M}}{4\cdot|G_E| + |G_{W_M}|} = \frac{4\cdot[0.02,-0.01] + [0.03,0.00,-0.02,0.01]}{4\cdot2 + 4}$.
Resulting fused gradient $\hat{g}$ has lower variance than either $g_E$ or $g_{W_M}$ alone.
The weight $\rho$ automatically down‑weights the noisy world‑model contribution when the measured residual is large.
**Figure 2.** WMRL corrects the world model rewards in two steps. Each row is one of the $m$ groups in a batch. Every group is graded by the world model into biased, noisy scores $\hat{r}$, and anchor groups are also graded by real execution. The monotone map $\hat{f}$ is fit on the score pairs (Online Debiasing) to remove the bias, and $g_E$ and $g_{WM}$ are fused by Inverse-Variance Denoising to lower the update variance.
Together, these components let WMRL keep the training loop fast while guaranteeing that bias and variance introduced by the world model are continuously corrected.
Convergence and Error Analysis
We bound how world‑model bias and noise affect RL convergence and show our corrections tighten the bound.
World‑model rewards inject a deterministic bias B and stochastic noise $\sigma$ into the gradient, inflating the error floor of RL training.
We view RL training as repeatedly moving toward the optimal score J* and ask how the world‑model’s bias and variance change the distance after T steps.
Compute c = 4M²/(n $V_E$) = 4·1/(10·1) = 0.4.
Compute `V_WM` = (1 + c $\sigma$²) $V_E$ = (1 + 0.4·0.0025)·1 ≈ 1.001.
Bias contribution = M² B² = 1·0.01 = 0.01.
Variance contribution = $\gamma$ `V_WM` = 0.5·1.001 ≈ 0.5005.
Geometric decay after T=5 steps: (1−$\gamma$$\mu$/4)^5 = (1−0.5·0.2/4)^5 = (1−0.025)^5 ≈ 0.88.
Bound on optimality gap ≈ 0.88·0.5 + 0.01 + 0.5005 ≈ 0.94.
The constant bias term (0.01) remains regardless of T, while the variance term dominates the error budget, illustrating why correcting both bias and noise is essential.
Why can’t we simply increase the number of training steps to eliminate the bias term?
The bias term adds a fixed offset to the gradient’s mean; it does not shrink with more iterations, so longer training cannot remove the error floor it creates.
Under Assumption 2, training on world‑model rewards yields J*−E[J($\theta_{T}$)] ≤ (1−$\gamma$$\mu$/4)^T $\Delta$₀ + O(M²B²) + O($\gamma$ `V_WM`).
Under Assumption 2 and the conditions of Appendix B.3, with high probability WMRL satisfies J*−E[J($\theta_{T}$)] ≤ (1−$\gamma$$\mu$/4)^T $\Delta$₀ + \tilde{O}\!\left(\frac{M²B²}{1+T/T₀}\right) + O\!\left(\frac{$\gamma$ $V_{WM}$}{1+$V_{WM}$/$V_E$}\right).
Thus, WMRL simultaneously contracts the bias floor and attenuates variance, matching the convergence of full‑cost RL while requiring only a tiny fraction of real executions.
Task Pool Construction
We revisit the theoretical lemmas that underpin the task‑pool design and show how each component contributes to reliable learning.
The paper’s core idea is to replace costly real‑world execution with a world model, then correct the model’s bias and noise using a thin stream of real data. This section examines the theoretical pieces that make the task‑pool robust.
The task pool is a curated set of simulated trajectories that the world model generates, together with occasional real‑world rollouts that keep the pool anchored to reality.
How does this task‑pool differ from a simple replay buffer that stores only real trajectories?
Unlike a pure replay buffer, the task pool actively injects model‑generated trajectories, which provide dense learning signals. The real rollouts are used only to correct the bias (the “$g_b$” term) and to control the variance (the “g_$\xi$” term), so the pool can remain small while still delivering a high‑quality gradient.
Remark 13 notes that restricting to linear combinations does not sacrifice optimality because, under Gaussian assumptions, the inverse‑variance weighting is also the minimum‑variance unbiased estimator among all possible functions of the two estimates.
Formal Proofs and Notation
Reference symbols and detailed proofs supporting the paper’s claims.
This appendix collects the notation used throughout the paper and provides the full mathematical proofs omitted from the main text.
This table lists symbols and their corresponding meanings used in the paper.
The proof sections (B.1–B.4) rigorously establish the convergence properties of the proposed world‑model‑based RL algorithm and the auxiliary lemmas used in the main theorems.
Experimental Configuration
Details on how the task pool is built, split, and curated for experiments.
We replace MLE‑Bench with MLE‑Dojo for three reasons. First, two original Kaggle competitions have been shut down, so their data and submissions are unavailable. Second, the medal‑based scoring system flattens at our model scales, masking improvements on most tasks. Third, many competitions predate current pre‑training corpora, introducing contamination risk.
The partition follows a deterministic rule fixed before training. Within each modality (tabular, text, image) we select the 20 smallest competitions, use the 15 smallest for training and reserve the 5 largest for held‑out evaluation, yielding 45 training and 14 held‑out tasks. Audio competitions and a task requiring an unavailable grader are excluded.
DSBench’s transfer suite provides the underlying Kaggle tasks; we keep the 60 that run end‑to‑end in our sandbox and discard 15 that are too large, lack a sample submission, or exceed column limits. The retained tasks span binary, multi‑class, regression, and other metric categories, and we verified no name overlap between training and held‑out sets.
**Table.** Categorization of datasets by modality: Tabular, Text, and Image.
**Table 1.** The transfer suite is the data-modeling split of DSBench, whose tasks are all drawn from Kaggle.
Prompting Strategy
All prompts, sandbox setup, and calibration details needed to reproduce the experiments.
This appendix records every prompt and sandbox configuration used for the AutoResearch and VLA agents, as well as the calibration procedures that keep the world‑model bias in check.
Training runs use GRPO with eight groups per step, eight trajectories per group, up to four interaction turns per trajectory, and a token budget of 4096 per turn; observations are truncated to 1024 tokens.
Sandbox environment packages (Table 7)
The sandbox runs each solution in an isolated Python 3.11 virtual environment on Linux with a single ~40 GB NVIDIA GPU; the environment contains exactly the packages listed above.
World‑model prompting is done without fine‑tuning: the prompt repeats the task description and the agent’s current solution, asks the model to simulate execution, and caps the context at 12 k tokens with a 1024‑token prediction budget.
Recalibration follows a monotone map: it remains the identity until 200 anchor pairs are collected, fits the map at that point, and refits after every 64 new pairs to track world‑model drift.
Reference disagreement $\hat{\eta}^2$ is normalized by the within‑group spread of true scores; a single warm‑up measurement yields a conversion factor $c = 1/\hat{\eta}^2 \approx 0.96$, confirming it is of order one.
Anchor weight is bounded to $[1, w_{\text{max}}]$ with $w_{\text{max}}=4$, preventing down‑weighting below predicted groups and limiting spikes that could let a few anchors dominate a batch.
VLA policy uses MiniVLA‑1B, a Prismatic‑style model with a Qwen2.5‑0.5B backbone and a vector‑quantized action head that encodes an eight‑step action chunk into seven discrete tokens.
AutoResearch agent system prompt (excerpt)
AutoResearch agent user prompt (excerpt)
The world model for the AutoResearch agent receives a system prompt describing its role as a precise execution‑and‑grading simulator and a user prompt containing the task overview and the agent’s code with line numbers.
World model system prompt (excerpt)
World model user prompt (file‑reference check excerpt)
VLA agent prompts are unchanged from the upstream Prismatic codebase: a single system turn with the stock Qwen chat template and a user turn containing the camera image and a lower‑cased instruction.
VLA agent system prompt (stock backbone)
VLA agent user prompt (excerpt)
The VLA world model (Robometer‑4B) receives only the instruction and eight uniformly sampled frames; it predicts per‑frame progress (float in [0, 1]) and a success probability, without any system prompt.