Reinforcement Learning for Machine Learning Engineering Agents
Authors: Sherry Yang
Joy He-Yueya
Percy Liang
Venue: ICLR 2026 Poster (2026)
Keywords: Machine learning engineering
language model agents
reinforcement learning
Primary area: reinforcement learning
TL;DR: We show that adapting a small language model agent with reinforcement learning using duration-aware gradients and environment instrumentation can outperform prompting a much larger, static LM agent for a specific ML engineering task.

Abstract: Machine learning engineering (MLE) has a clear objective: Given an MLE task and a verifier (e.g., performance on some held-out data), what is the most effective way to utilize compute to achieve the best performance for the given task? Existing language model (LM) agents rely on prompting frontier LMs and accumulating experience non-parametrically by storing and retrieving experience through agent scaffolds and test-time compute. In this paper, we show that in environments such as MLE where a good verifier is available, adapting the LM parameters through gradient updates can be more effective in utilizing compute and agent’s experience. Specifically, we show that agents backed by weaker models that improve via reinforcement learning (RL) can eventually outperform agents backed by much larger, but static models for a given MLE task. We identify two major challenges with RL in this setting. First, actions can take a variable amount of time (e.g., executing code for different solutions), which leads to asynchronous policy gradient updates that favor faster but suboptimal solutions. We propose duration-aware gradient updates in a distributed asynchronous RL framework to amplify high-cost but high-reward actions. Second, using performance on the held-out data as a reward for MLE provides limited feedback. A program that’s nearly correct is treated the same as one that fails entirely (e.g., during data loading). We propose environment instrumentation to offer verifiable partial credit, using a separate, static language model to insert print statement to an existing program. Our experiments suggest that a small LM (Qwen2.5-3B) adapted with RL, when given enough compute, can solve an MLE task better than prompting a frontier model (Claude-3.5-Sonnet) with the state-of-the-art agent scaffold (AIDE) by an average of 22% across 12 Kaggle tasks.

Published as a conference paper at ICLR 2026
REINFORCEMENT LEARNING FOR MACHINE LEARNING ENGINEERING AGENTS
Sherry Yang." 1.2 Joy He-Yueya? Perey Liang?
NNww York University "Stanford University
sherzyyangenyu.edi
ABSTRACT
Machine learning engineering (MLE) has clear objective: Given an MLE task and a verifier (e.g., performance on some held-out data), what the most effective way to utilize compute to achieve the best performance for the given task? Existing language model (LM) agents rely on prompting frontier LMs and accumulating experience non-parametrically by storing and retrieving experience through agent scaffolds and test-time compute. In this paper, we show that in environments such MLE where good verifier is available, adapting the LM parameters through gradient updates can be more effective in utilizing compute and agent's s experience. Specifically,
show that agents backed by weaker models that improve via reinforcement learning (RI.) can eventually outperform agents backed by much larger, but static models for given MLE task. We identify two major challenges with RL in this setting. First, actions can take a variable amount of time (e.g., executing code for different solutions), which leads asynchronous policy gradient updates that favor faster but suboptimal solutions. We propose durution-aware gradient updates in a distributed asynchronous RL framework to amplify high-cost but high-reward actions. Second, using performance on the held-out data as a reward for MLE provides limited feedback. A program that's s nearly correct is treated the same as that fails entirely (c.g.. during data loading). We propose environment instrumentation to offer verifiable partial credit, using separate, static language model to insert print statement an existing program. Our experiments suggest that a small LM (Qwen2.5-3B) adapted with RL, when given enough compute, can solve MLE task better than prompting frontier model (Claude-3.5-Sonl
with the state-of-the-art agent scaffold (AIDE) by an average of 22% across Kaggle tasks'.
INTRODUCTION
Machine learning engineering (MLE) aims to answer the question: Given an MLE task, what the most effective way to utilize compute to achieve the best performance on the held-out data? Existing MLE agents rely on prompting frontier LMs with agent scaffolds (Chan et al., 2024), accumulating experience non- oonpparametricllyy (e.g.. by storing and retrieving previous experience) (Jiang et al., 2025). While this approach of scaling up test-time compute (Snell et al., 2024; Wu et al., 2024) can allow an agent to search for better solutions, the agent's 's fundamental behavior does not change drastically without gradient updates, which wastes valuable experience from running costly ML experiments. As shown in Figure 1, running the best agent scaffold according to ML.BBench (Chan et al., 2024) with Claude3.5 ---ooneet for days leads to only slightly better best solutions.
A natural approach to improving an LM agent given past experience for an MLE task is to adapt its parameters through gradient updates using reinforcement learning (RL) (Sutton et al., 1998). However, agentic settings pose additional challenges to RL. First, action exccution of MLE agent can take a variable amount of time (e.g.. training different ML models), which leads to asynchronous policy gradient updates that favor faster but suboptimal solutions. To overcome this challenge, propose dturation-warr gradient updates during distributed asynchronous RL training to balance gradient updates faster and slower actions. With duration-aware updates, we observe that an agent stops favoring faster solutions, achieving better performance in the long run.
"oode available http:///ghhua.commseerryyhhie

---

Published as a conference paper at ICLR 2026
RL training of small model
Gaddeett update
Mox Soore (3 cowebt))
-0.5
Prompting large model
No gruthient update
1.0
Ower2. 5 30
5.Sunne Agunt Scuffold
-2.0
10
50
40
50
Time (tterss)
Figure I: Performing gradient update with RL on Qwen2.5-3B (blue) is more effective improving the best task performance than prompting Claude3.5-Sonnet with the best agent scaffold (red) on the "leaf-classification" task from MLEBench (Chan et al., 2024).
The second challenge of RL for MLE agents that, while performance on the held-out data serves as a verifiable reward, they offer limited feedback, treating a nearly correct program (e.g., failed during writing solution to the correct location) the same as one that fails entirely (e.g., during data loading). Such limited feedback hinders the progress of RI,, and further fails to capture if an agent uses MI. at all to solve problem that meant to be solved by MI.. address limited feedback, we propose environment instrumentation to offer partial credit to intermediate steps of completing an ML task (e.g., loading the data, building and training a model). We implement environment instrumentation by using static copy of the original LM to insert print statements in the code generated by the agent, the execution of which provides partial credit. We observe that partial credit can gradually guide the agent away from making trivial mistakes (e.g., import errors, failures to load data) and towards improving MI. techniques (e.g., feature engineering and hyperparameter choices).
Across set of 12 challenging MLE tasks, show that adapting the parameter of small agent (Qwen2.5-3B) through RL can eventually outperform prompting a frontier model Cllaude-..-Soonee with the state-of- -frte--rtr agent scaffold (AIDE) for 8/12 tasks, achieving average improvement over the frontier model by 22%. Our results suggest that future MLE agents should learn to balance the compute spent across inference, interaction (action execution), and performing gradient updates, as opposed to only focusing on scaling inference-time compute and interactions.
2 BACKGROUND
In this section, we provide relevant notations and define key learning objectives. We further discuss a few challenges of running standard RL algorithms in agentic settings.
MLE Agent in Markov Decision Process (MDP). We consider a Markov Decision Process (MDP) (Puterman, 2014) represented by a tuple (S, A, R, P.u). consisting of a state space S, an action space A, reward function R S x A + R, state transition probability function P S x A ASS), and an initial state distribution μ E (SS) A policy π S + A(A) interacts with the environment, starting from initial state so μ.. At each interactive step k 2 0, an action d ~ т (**)) is sampled from the policy and applied to the environment. The environment then transitions into the next state 8k+1 P(-(sk, ax) while scalar reward R(sk, ag) is produced. Reinforcement Learning (RI. aims to find policy π that maximizes the expected future rewards:
K
((T) E.
Σ R(Sk, ax)
LA=
where K is the total number of steps. Standard RL algorithms such as policy gradient (Williams, 1992; Schulman et al., 2017) can be applied to learn the policy update rule by estimating
K
V = Er.y. P.P
V log (s(ax|sk) A(sk, ax)
(2)
k=
where A(sk, ag) ) is some advantage function which can be separately estimated (c.g., by Monte-Carlo returns from (Williams, 1992)).
2

---

Published as a conference paper at ICLR 2026
Exceorico time of angentggeeennn pograns
120
110
GridSearch (0:52)
2
Import pandas as pd
100
Lest_dr readedd "/input/test.cowv
90
80
70
Randoma Forest Classifier (0.65)
5
def jaccard st
- set (strl, (1 spplit
60
set (sts2. (1 sspiit
60
= a.intersestion (b)
40
recurn float (1en / (len(a) len
80
Lineur logistic regression (0.64)
len(c))
20
for
test_ itterrows
10
phrese = nax (zow("Lext "eett]_aa .spli:() key- lambda x: jacsard scw["texs"?
10
1
22
24
training steps
Figure 2: Execution time of agent-generated pro- grams (averaged across 128 samples) decreases dras- tically RL training progresses. Plot produced by running distributed (Sheng et al., 2024) on the random-acte aertommacts-oo -ot za task (text classifi- cation). The final solution converges to a fast but suboptimal solution (achieving score 0.64) of using linear logistic regression that takes less than second execute, as opposed to other solutions with better performance (0.65) but takes longer to run.
Figure 3: Suboptimal convergence due to limited feedback. ing sentiment-relevatt phrases from tweets (t (tweet sentt ieent-exx racti the agent converged
a task extract-
suboptimal solution of directly coding the Jaccard similarity and search for the best phrase in the test input, bypassing ML completely. This demonstrates how sparse rewards can lead to an agent exploiting evaluation metrics rather than learning desired behaviors.
In the MLE agent setting, S captures input to the agent, including problem description, datasets, and any experiment history. A captures solutions generated by an agent, including high-level plans (e.g., which family of models use) and low-level code. P captures any potential output from the enviromment during action execution (e.g., error messages, printed training losses). Note that P is stochastic, as the outcome of training ML models (initialized with random weights) is non- deterministic, and good MLE agent should learn to perform well despite such stochasticity. An agent can generate solution = 1) followed by subsequent debugging or improving steps (k >> 1). R. captures rewards from the environment, such as performance of the MI solution on the held-out data.
Challenges of RL for MILE Agents. For efficiency reasons (Thrun, 1992; Kakade, 2003), many RL training frameworks implement asynchronous distributed setup where multiple "actors" can interact with their own instances of the environment simultaneously, gathering experiences which then sent to "learner" for distributed RL updates (Liang et al., 2018; Hoffman al., 2020). In agentic settings such as ML engineering, each action may take a variable amount time to execute. As result, running distributed RL training favors faster actions (slower actions might often time out). Moreover, time-consuming actions are sampled less frequently in a distributed training framework, leading to an uneven number of gradient updates for faster and slower actions. As shown in Figure naively running a distributed RL framework (Sheng et al., 2024) on a text classification task leads to the agent only generating quick solutions that barely take any time to execute.
Another challenge of RL for MLE agents is the limited feedback for intermediate progress. While natural reward, does not distinguish between solu-
performance on the held-out data tion failing to load data and one that is nearly correct. Furthermore, this lack of intermedi- reward can lead an agent not using ML to solve problem at all. For instance, in the tweet-e - sent iment rextraction ion task where the agent needs to extract coding the Jaceard phrases from tweets, the agent converged to suboptimal approach of directly similarity evaluation function and search the test input for the best phrase (as shown in Figure 3), bypassing ML completely.
3 RL FOR MLE AGENT
In this section, propose duration-aware gradient updates (Section 3.1) and environment instrumen- tation (Section 3.2) to overcome the aforementioned challenges of applying RL to MLE agents in Section 2. Moreover, the agent can further improve previously generated solution, which can be further enforced using RL (Section 3.3). See Algorithm in Appendix A.l for the training loop of the RL agent.
3

---

Published as a conference paper at ICLR 2026
Regex matching, e.g.. rearcach scaach""ooodd data", out)
At Vo Vogogagl Ro (a(a
Partial Credit Extraction
1.M Inputs
LM Samples
Insert print statements, c.g.. print( "lond"d data")
Solve ML problem.
train.csv, test.csv.
Policy
plan 1 + code 1 (Ar_1)
plan 2 code 2 (Ar_2)
Environment Instrumentation
Code Execution
Improve previous
LM Agent
solution.
plan code (Ac.1)
Figure 4: Proposed framework overview. Duration-aware gradient updates re-weights the policy gradient updates according to the execution duration of an action. Environment instrumentation inserts print statements using a static LM, the exccution output can then be extracted for partial credit. The agent can further asked to improve previous solution, the response can further be enforced with RL.
3.1 DURATION- AWARE GRADIENT UPDATES FOR VARIABLE- TIME ACTION EXECUTION
A naive solution to the aforementioned problem of variable duration actions in Section 2 is to wait for all actions finish executing before performing any policy gradient updates, but this does not utilize resources well and is not scalable as training a model for a hard ML problem can take days.
The Issue with Mariable-Duratiin Action Execution. We first provide a precise illustration of the issue with suriable-time action execution in distributed RL training. Define n, and ny as the number of samples of actions x and y collected in time T. Denote Åg and Ag as the average advantage estimates for actions x and In a fixed training period of length T, we would collect approximately: T Tx)())
π {u]s)- -T
T2
samples of action x,
ny At,
samples of action
Ata
Where (x)s)) and (t)s) are the probabilities of selecting actions and y under the current policy. The total gradient contribution for each action would be:
Ge (xrs)) AC
Gy T(yls) T V logg Åy
V log (x(ss)) A.
Aly
Note that G2 and G. are divided by At, and A,y, meaning faster actions (smaller At) contribute proportionally more to the gradient.
Duration-Aware Gradient Updates. To counter the frequency bias above, we propose weight each gradient update by the action duration, which gives
G² π (z]s)) T Ve log To(=]s) As,
=)=x(yls)) T Vo (ogomeylss) Ay
With duration weighting, the Al terms cancel out, leaving each action's 's contribution to the gradient proportional only its policy probability and advantage, that actions with longer durations receive fair consideration in policy updates despite generating fewer samples in the same time period. Generalizing from this toy example to the continuous case with arbitrary action durations, we arrive at our duration-awen aware policy gradient update rule:
its execution frequency. This ensures
(oJ(TO) = ET
Ate Vo log To (axls))) Á(Sı,
ax)
Where At, is the execution duration of action as taken state sy. This formulation ensures that in expectation, the contribution of each action to the policy gradient is proportional to its true value, regardless of how frequently it is sampled due varying execution times. In practice, we rescale A/e by the average execution time the batch to avoid overly large gradient updates. See Appendix A,2 for the proof of the unbiased policy gradient estimate with duration-ane -ware aware gradient in a distributed RL setting.
3.2 ENVIRONMENT INSTRUMENTATION FOR PARTIAL CREDIT
Environment Instrumentation. To overcome the challenge of sparse reward in MLE, we propose to introduce partial credit so that generated programs that fail in the beginning (e.g., during data loading) will receive less partial credit than programs that is almost correct (e.g.. failed at saving output to the correct location). To avoid making too much assumption about how the agent should solve a problem, we assign partial credit only based whether solution completed high-level procedures including importing libraries, loading data, building ML model, training the model, and running the model on the held-out data. We propose instrument the code generated by the agent by inserting print statements into the program generated by the agent to track execution progress. The terminal output will be parsed through regex matching
another static copy of the original LM to
4

---

Published as a conference paper at ICLR 2026
Action: plan and code generuted by the agent
Environment Instrumentation prompt
Free
iaport Aer ansooeetiiiit
Insert print statements python program after the conesponding code:
rssc a random fores classifier walve this will first load data from 'inputt and perfocm of idf on
Insernte pacoqast
"rquest sexr, edit, aware will fimally implement
print"inppeial packages")
loaded data")
test pe.rea Jaan compt//ett..prr
randon forest and save
subenision.co cay.
trist vectoriaer Triffaacterteer
regurtt tort 2013 parcis
Environment LM
Tl aedeerratiaaattion R
printe Verivee wakery
firt
cart, Aattt recuent undli
Code Execution
stear1)
classstion desettrenttleettiie
Irint traiied sedet)
Crain Caseeeiittrrgeeloooooeeeeeel
"e-search"" packages", out! -L rb, seseeeercd asded rata", outi ras Si.sacckhhrraiee necel", out)
Frepirt seeetseees rite
Frepare apeeaases fi1
Treautiier, rectiet si y.rret
inantertoe.ne to, StA"". hra.i///aaeteeeee.ee... steee racod
Extruct partial credit
Figure 5: Environment instrumentation overview. Another copy of the small LM (Qwen2.5-3B) is prompted to insert print statement into the code generated by the agent. After code execution, output from the terminal then parsed to assign partial credit by regex matching.
to provide partial credit based whether expected print statements (e.g., "print(loaded data)") are executed, as shown in Figure 5.
Preventing Partial Credit Hacking. It is important to note that environment instrumentation is done by a separate, static model. If we have the same model being optimized by RL generate the print statements, we would very likely see reward hacking. By freezing the model used for instrumentation, we have observed that it largely mitigates partial credit hacking. However, instrumentation from a static model is still possible to be hacked. To further prevent partial credit hacking, we use different scales for partial credit and final reward (-10 for the worst solution, 0.1 for each of the 7 candidate print statements. If a generated programs runs without error, the submission (label for the held- out data) is graded by the grader of the environment, and true task performance is used as reward (generally between- -1 and 1). As result, the model will see much higher reward for producing valid ML solutions than to hack the print statements.
3.3 MULTI- RL WITH IMPRVVENTTT PROMPT
So far, we have discussed the setting where an agent is directly asked to generate plans and code solutions for solving MLE tasks. Next, we further explore whether we can directly instruct the agent to improve a previously generated solution. Specifically, we sample from two sets of prompts (with equal probability) to solve the problem from scratch and improve solution generated by the previous step. We illustrate the two types of prompts in Figure 4. In the case of improving a previous solution, output of the terminal is given to the agent which includes information such training and test accuracy (from environment instrumentation introduced in Section 3.2). We have also experimented with giving failed executions to the agent to self debug, but have noticed limited self-debugging abilities in small models. At test time, we both generate solutions from scratch and run the agent again to improve the generated solutions, and take the maximum between the two solutions (with and without explicit improvement).
4 EXPERIMENTS
In this section, we evaluate our proposed improvements to adapting LM agent to MLE with RL. We first discuss the evaluation setup and implementation details in Section 4.1. We then present the main evaluation results in Section 4.2, followed by ablation studies in Section 4.3.
4.1 EVALUATION SETUP AND IMPLEMENTATION DETAILS
Evaluation Setup. We perform evaluation on 12 MLE tasks from Chan et al. (2024) spanning vision, language, and tabular inputs. These tasks are selected based on whether a small model (Qwen2.5-3B) can produce a submission with valid format, because we are only interested in improving MLE, not in the general instruction following and coding abilities (which is what is required to produce submission with the valid format). We use the grader from Chan et al. (2024) to grade the final 128 samples after RL converges and measure both the mean and the maximum performance for each run (see additional studies on the effect of sample size in Appendix D.3). We
5

---

Published as a conference paper at ICLR 2026
Table 1: Comparing RL of small model to prompting large models across 12 tasks from MLEBench. RL results are best scores among 128 samples after RL has converged. Baseline results are from runs Chan et al. (2024), produced by prompting frontier models using AIDE agent scaffolds and continuing running for 24 or 100 hours. Numbers shown are mean and standard error across 3 runs. All except for the last column use the AIDE agent scaffold. 1 denotes the higher the score the better. N/A denotes no valid submissions were available. RL of small model achieves the best final performance on 8 out of 12 tasks.
Tasks
Qwea2.5-3B Llama3. 1-405B Clandes, Sonn GGPT-fo-00trs
Qwen2..5335 5:3B RL
detecting insults in social commentiryy (1)
0.870 +/ 0.009
N/A
NIA
N/A
0.895 +/. 0.001
learning agency- Saas--hh--aaaan tah-automatedd-ee 2 (1)
0.331 +/. 0.018
0.777 +/ 0.002
0.794 +/- 0.008
0.759 +/- 0.002
0.746 +/. 0.002
random-acts- offf-------- (1)
0,589 0.004
0.619 +/- 0.007
0.627
0/00:
0.638
0.005
0.663 +/- 0.011
(t)
0,027 +/- 0,018
0,939 +/- N/A
0,448
0.251
0.283 +/- 0.005
0.596 +/- 0.002
labular playground seriee-eeeeeeee may
2022 may-2022
(t)
0.787 +/ 0.020
+/- 0.126
0,883
0.002
0.000
tabular playground series dec 2021 (?)
0.827 +/- 0.044
0,771 +/ 0.188
0.645
0/.0.15
0.957 +/- 0.000
0.951 +/- 0.000
us patent phrase-to teaann------------ee matcciintnn (t)
0,065 +/- 0,000
N/A
0.805 +/- 0.006
0.588 +/- 0.015
0.527 +/- 0.003
plant pathology 2020 2020.fgve7 (f)
0.628 +/ 0,058
0.968
0.005
0/990.+
0:970 +/ 0.001
0.970 +/ 0.004
leaf-el classsssifiitioon (L)
0.884 +/- 0.016
6.747
+/- 5.398
0.436 0/-
0.846 +/- 0.029
0.124 +/- 0.000
nomad2018-ppo 2ere000--dddd--raaaeee
sulcc-ranspaccnt-oooduuee (4)
0.178 0,045
0,166 +/- 0.103
0.083 0/000020
0.072 0.003
0.059 0.000
(4)
0.596 +/- 0.053
0.487
+/- 0.020
0,701 +/- 0.186
0,546 +/- 0.004
0.404 +/- 0.011
Imays chatbot arena (t)
11.48 +/ 0.002
1.269 +/ 0.051
2,211 +/ 1.51 +/ 0.035
1,081 0.002
use the scores achieved by different frontier LMs and different agent scaffolds from the original runs of MLEBench baselines. The baselines prompt frontier models with the state-of-the-art AIDE agent scaffold, which organizes experience in a tree structure and saves the best solution seen so far for evaluation. In evaluating against different agent scaffolds, we use the results from the GP based agent rumning 24 hours using two additional agent scaffolds, OpenHands ( (Wang et al., 2024) and MLAgentBench (MLAB) (Huang et al., 2023b) (which has outperformed LangChain (Chase, 2022) and AutoGPT (Significant Gravitas)). To further understand the improvement progress of RL. and prompting, re-run the set of MI ML.EBenc experiments using -.aooe--.5-Sonet and AIDE agent scaffolding, while grading the intermediate best saved solutions, and compare that to intermediate solutions during RL, both across three runs.
Implementation Details. To implement RL for the Qwen model, we build on top of the distributed RL training framework in Sheng et al. (2024). We implement a set of distributed sandboxed code execution environments similar to Chan et al. (2024), where code execution takes place inside of the RL training loop as part of the reward function implementation. To implement environment instrumentation, we load separate copy of the original Qwen2.5-3B model (without performing any gradient updates on it) and ask the model to insert print statements before executing the code. The prompt for environment instrumentation can be found in Appendix C.3. To assign partial credit, we use reward- -10 to denote solutions that fail completely (e.g., no plans or code, fail to import packages), and add 0.1 per regex match in the terminal output. If the solution is valid (according to the grader), we use the actual score from the grader as reward. We further experimented with normalizing the reward to a particular range (0 to I) but did not observe significant difference. For tasks where the lower scores are better, we flip the signs of the scores to use as rewards. We use the Proximal Policy Gradient (PPO) (Schulman et al., 2017) algorithm (with modification of duration-awar gradient) to train the Qwen2.5-3B model for each task until reward convergence, which generally took 1-3 days depending on the task using 8 A A100-40GiB GPUs. We limit the input and output length to 1024 tokens. The model is trained using a batch size of 128 and learning rate of 1c-5 (See complete set of hyperparameters in Appendix B.
4.2 EVALUATION RESULTS ON MLEBENCH
Comparing against Different Frontier Models. We report the mean and standard error across three runs of RL or prompting a frontier model using the AIDE scaffold in Table 1. Since the AIDE scaffold saves the best solution found through prompting, we also report the maximum among the 128 samples (See the mean scores across the 128 samples during RL in Figure 10 in Appendix D. I.... Qwen2.5-3B with RL. outperforms prompting a frontier L.M on out of the 12 tasks, and achieves an average of 22% improvement (measured by improvements or degradation over the baselines) over prompting Tlaude3.5-Sonnee models, we still observe significant improvement running RL over prompting the Qwen model with the AIDE agent scaffolding (As shown in Column 2 of Table 1). We note that running AIDE for longer hours (e.g., GPT-40 100 hours as opposed to the default 24 hours for other columns) did not lead to significantly better performance, indicating that solutions other than prompting a large frontier LM is required to effectively achieve seff-improvemmnnt
For the tasks where Qwen could not outperform the frontier

---

Published as a conference paper at ICLR 2026
Generated plan
Generated code
To improve the previous sulution adhiave lower log loss on lest ad, we cmn improve "Fexuure Engineering** We can cranee ndlitional festurs from the prompe and ropprses, auch asthe length of the the presence cortain keywords, the number winds.
Feature encineering
response-a_tooyee a_length]] trasn
rrespoose an
treanddt respense
tratn_ari response dirf'l Strain arl response a length'1
resposs, leggu)))
Additional Fecture Engineering the word count di fleeenos and evergg word leugt differnnce between the two responses are exctted to eapture more infommatios about the resporees.
use additional
Adtitional festures
train
wore dittr trate resoonse 2il 2'1-1polyc
lursta xt lentx.splatii splitt))) - tf ('response_ bi Lespplyt lanses st tents.plaatt))) retpreseenee
peespose
lamoda rp.scant (lniwardd word x-salitt) 11 crais_afi respansee (TTaapl
Lantes x: no. want Lentierd) for verd in is s-spsst))))
Figure 6: Qualitative examples of improvements proposed by the agent during RL. [Top] earlier improvement proposed by the agent using difference between response length as features for preference prediction. [Bottom] later improvements proposed by the agent using additional features such as word count and average word length difference as features.
Table 2: Comparing RL to different agent scaffolds. Adapting a small model with RI. outperforms prompting GPT-4o with different agent scaffolds on 9 out of the 12 tasks. Results for agent scaffolds are taken from Chan et al. (2024). Numbers show mean and standard error of the final performance according to the grader in Chan al. (2024).
Tasks
GPT-4o AIDE
GPT-4o OpenHands
GPT-4o MLAB
Qwen2.5-3B RL
detecting-insaffe wiwpangiiiinnnnnnnniiiieeeee wtweiii----ciiioiiiiii commentary (f)
NaN
0,867 +/- 0.017
0,749 4/- 0.039
0.895 a/- 0.001
learning-agency- saaiiig-ggnyy-l---auueooe
-aw---atonaatdd-sss---oooiiiiiie (f)
0,720 +/- 0.031
0.681 +/- 0.010
0.533 +/- 0.080
0.746 +/- 0.002
random- aftss-o---pipaaa pizzi (t)
0.645 +/- 0.009
0.591 +/- 0,048
0.520 +/- 0.013
0.663 +/- 0.011
0.294 +/- 0.032
0.415 +/- 0.008
0.158 +/- 0.057
0.596 +/- 0.002
aabular-payy srooundd-eriis--ee
(?)
0 0.012
0,882
+/-
0.913 0.000
lahular playgr playground series dec:2021 2021 (1)
0.957 +/- 0.002
0.957 +/ 0.000
0.828 +/ 0,118
0.951 +/ 0.000
us-paten patent- phrase-to phrase matching (f)
0.756 0.019
0.366 0.019
NaN
0.527 0.003
plant- plant-patholes pathology- bathology-2020-fggve 2020-fgve7 (f)
0.980 +/- 0.002
0.680 +/- 0.113
0,735 4/- 0.052
0.970 +/- 0.004
leaf-clasifliationn (L)
0,656 +/- 0.070
0,902 +/- 0,018
4.383 +/- 2.270
0.124 +/- 0,000
nomad2018 predict transparent
cot.ppparrntccccccteseessss conductors (4)
0.144 +/- 0.031
0.183 +/ 0.120
0.294 +/- 0,126
0.059 +/- 0,000
booky--outoor-idnntiiialldd (1)
0,576 +/- 0.071
0,582 +/- 0.020
0.992 +/- 0.463
0.404 +/- 0.011
lnssys-chatbut arena (4)
1.323 0.147
1.131 0.019
+/- 4.509
1.081 0.002
In Figure 6, we provide example solutions that the Qwen agent came up with during RL in solving task. This task requires the agent to come up with ML code and train a model to predict which responses generated by LMs a user will prefer. The Qwen agent is able to come up with various different feature engineering choices such as using the difference in response length, word count, and average word length as additional features.
the sesyss-chatottaaee
Comparing against Different Agent Scaffolds. We now compare running RL on Qwen2.5-3B against running different agent scaffolds on GPT-4o. Table 2 shows that Qwen2.5-3B with RL outperforms prompting LMs with various agent scaffolds on 9 out of the 12 tasks, achieving an average improvement of 17.7% over the best scaffold for each tasks. We notice that the performance of prompting a frontier model does vary for each task across different agent scaffolds. For instance, AIDE achieved no valid submissions across 3 runs (score NaN) for Llama3.1, Claude3.5 -Sonnet, weettiingnnnnsssee
and GPT-4o achieve valid solutions on the same task. Nevertheless, AIDE generally works better in achieving higher task performance compared to other agent scaffolds, suggesting that RL is reliable way to improve performance that is agnostic to different choices of scaffolds.
while other scaffolds can
Performance Improvement over Time. In Figure 7, we compare the max perfor- mance (across 128 samples) aggregated over time of training Qwen2.5-3B with RL against running AIDE agent scaffold using Claude-3.5-Somnet. many tasks such as sweet-sentin iennttme-ttnee wetennttextttiiiiii model initially achieves much better performance than the small model. However, as RL goes on, performance of the smaller model improves more with gradient updates, eventually exceeding prompting large model.
We observe that for lawtllllaaauttaaae aut -ssay-y-eeessyssoorrrr and random acts-offpil a--of-pizza, prompting the large
4.3 ABLATION STUDIES
We now present ablation studies on auration-aware gradient updates, environment instrumentation, and explicit self-improvement prompt.
Effect of Duration-Award Gradient. We plot the average execution time (across 128 samples) during RL with and without duration-awart gradient in Figure 8. We see the overall runtime still decreases as RL progresses, since tigh-performing programs generally tend to be fast. However, with duration-aware gradient, the agent is able to find better solutions that take longer to execute (e.g.,
7

---

Published as a conference paper at ICLR 2026
www.
(eee) S0p EJ
(m022/00 &
Gmei 500 E
speily cedipralaar Certeiaaaa
0002
130
Orita 200
0.13
Figure 7: The best scores achieved by the agent across time comparing prompting a large model to adapting a small model with RL. A small model running RL starts off with low scores for many tasks, but eventually outperforms prompting large model.
Exccution of aaneeratedd programs (6)
No DAG
Duration Aware Gradient
model Placklact ("erid y Trinfeector tokent ze x.splin())), Gradientaboatinng lessifientt)
120
More expensive
110
II
fintrri tects, trait Irequester recetved pirrs
prot_ wotel.prodiet tost texts) It, 11
Score 0.67
30-
tridf vectorizer TridfVector lac tstp wordse Lish"))")
X. train triet vertorieerr,, fit transtorat trrfforrttttttttiii request sontt_ddtt..orrr II
teidf_vector www.orrreer.trrneeeee tranerttaasffttteeeee
request text _eiterr
Train Topistic regression wodel.
model segistisejresssoon
reqq rrquesterr_eceevvved reeeeeeeee__ppeee
10
steps
prebasititics on the probabilittes
Score
Figure 8: Duration-aware gradient enables the agent to explore more expensive but high-return actions by coming with more expensive solutions such as gradient boosting, which achieves higher score than linear logistic regression. The RL agent still tends to find faster exccuting solutions over time.
gradient boosting), whereas without duration-aware gradient, the agent quickly converges to fast but suboptimal solution (e.g., linear logistic regression). Nevertheless, tends to find faster executing solutions, as the average execution time still tend to decrease over time. We found that best solutions from PPO+DAG achieve an average of 3.85% improvement over vanilla PPO.
found that the RI. agent still
Effect of Environment Instrumentation In Figure 9, we show the average task scores (across 128 samples) during RL with and without environment instrumentation. The task scores are -10 if the solutions are invalid and the actual scores otherwise (partial credits from environment instrumentation are omitted from the plot but is included in the actual reward RL optimizes). We observe that environment instrumentation leads to faster growing and faster converging average scores. The high- variance in pl ant -httho logy-2020-20 not being able to produce any valid solution due to sparse reward, which we observe more frequently when environment intrumentation is absent. We found best solutions from PPO+EI achieve 22.06% % improvement over vanilla PPO.
(right most subplot) was due to one RI. training run
No Envirenment Instrumentaaion
Environmentt Instrumentation
tweet-senti www.eeeeeeettttttttoooo learning-agenccy- httpa/---gggg----a-atttt-tt wwtp---bbbttttt---oooooo puttraate--ssyy--orrm-
plart-pathel -2922
2-
0-
10
E324
120
Aep
training steps
Figure 9: Environment instrumentation ablation. Plots show the mcan task scores (excluding the partial credit from environment instrumentation) across 128 samples for 3 example tasks across RL training steps. Environment instrumentation improves RL and enables faster convergence.
8

---

Published as a conference paper at ICLR 2026
Effect of Explicit Self-Improvement Prompt. Next, we compare explicitly asking the agent to improve a previous solution (50% of the time during RL) to only having the agent solve the task from scratch. Asking the model to improve a previous solution leads to better final performance for 10 out of 12 tasks and achieves an average improvement of 8% over coming up with solution from scratch, suggesting that RL can simultaneously improve both initial solution generation and improve the ability to improve previously solution. See the complete performance with and without self-improvement prompt in Table 4 of Appendix D.2.
RELATED WORK
ML engineering agents. Many recent work has emerged for building LM agents that can solve
machine learning benchmarks (Huang et al., 2023b; Tang et al., 2023; Chan et al., 2024; Li et al., 2024; Zhang et al., 2025), data science tasks (Grosnit et al., 2024; Bendinelli et al., 2025; Pricope, 2025), or help with other aspects of MI. such as data preprocessing and hyperparameter optimization (Zhang et al.; Liu et al.. 2024; Gu et al., 2024). Most existing work in this space has focused on prompting large frontier LMs as opposed to performing gradient updates. Existing work has used various agent scaffolds such as LangChain (Chase, 2022), AutoGPT (Significant Gravitas), OpenHands (Wang et al., 2024), and AIDE (Jiang al., 2025) for in-context learning, and further try to improve agent performance by heuristic- based search during inference time (Liang et al., 2025). While LM agents perform better with these scaffolds, they still face the challenge of achieving improvement reliably from prompting (Huang et al., 2023a; Errica et al., 2024). We focus on RL training of smaller models instead of prompting large models.
RL for LMs. Since the development of both policy and value based RL algorithms (Williams, 1992; Sutton et al., 1999; Kakade, 2001; Schulman et al., 50 2017; Watkins & Dayan, 1992) extensively took places in simulated environment such as MuJoCo (Todorov et al., 2012) and Atari (Bellemare et al.,, 2013), many RL optimization frameworks (Hoffman et al., 2020; Hafner et al., 2017) make the implicit assumption is that environment interactions take up a constant amount of time. More recently, RI. has been used extensively aligning LMs to human preferences (Ouyang et al., 2022; Rafailov et al., 2023; Christiano et al., 2023; Ziegler et al., 2020), reasoning (Lee et al., 2023), and solving math (Dang & Ngo, 2025) and coding (Wei et al., 2025) problems. However, this assumption persists, as rewards are often produced by reward model (Ouyang et al., 2022) or verifiable answers to math or coding problems (Guo et al., 2025; Wei et al., 2025). As a result, the problem of variable-tima time action execution has not been extensively studied. However, this problem is highly relevant in practical agentic systems such as ML engineering. As RL being extended to a broader array of agentic applications, deriving optimization frameworks that take into account the time an action takes is essential. Mcanwhile, directly applying existing RL training frameworks developed for simulation settings, math, and reasoning, such as Hoffman al. (2020); Sheng al. (2024), results in poor agent performance.
RL for agentic systems and interactive tasks. Existing work has studied RL for agentic settings solving multi-step interactive tasks such as operating unix terminal (Liu et al., 2023), booking flights (Snell al., 2022b), controlling devices (Bai et al., 2024), negotiating price ( (Verma et al., 2022), navigating through the web (Zhou et al., 2024), and playing language-based based games (Narasimhan et al., 2015; Snell et al., 2022a). However, most of these settings still neglect the time it takes to execute actions, and mostly leverage gamma discounting (Sutton et al., 1998) to balance the influence of future rewards and immediate rewards. These settings do not consider each turn taking a different amount of time. Additionally, sparse reward has been challenging in many agentic settings. Existing work has leveraged LMs/VLLMs as process reward (Choudhury. 2025; Zhang et al., 2024; Mahan et al., 2024) to provide dense reward signals for policy evaluation and RL training (Pan et al., 2024; Venuto et al., 2024; Bai et al., 2024). However, directly using LM as reward functions can be unreliable (Son al., 2024; Singhal et al.). We tackle sparse reward by having LM insert verifiable print statements as form of reliable execution feedback to improve the agent through RL.
6 CONCLUSION
We have shown that adapting LM through RL is better than prompting a frontier model with non- parametric memory for given MLE task, even with much smaller model. We have also shown that reweighting policy gradient updates based on action duration can overcome variable-duration interactions, while using an LM to perform reward instrumentation through code can mitigate sparse rewards. These findings suggest future work on MLE agents to focus more on adapting model parameters with verifiable rewards, and offer tools for resolving challenges of RL. in MLE settings.

---

Published as a conference paper at ICLR 2026
Due to limited compute and small base models available, only demonstrated the effectiveness of the proposed methods on 12 MLE tasks, where the small model was not bottlenecked by basic instruction following capabilities. Nevertheless, this research is the first to show what RL training could achieve for MLE even with a small model and limited compute. We expect our findings to hold for more tasks with larger and better base model.
REFERENCES
Hao Bai, Yifei Zhou, Jiayi Pan, Mert Cemri, Alane Suhr, Sergey Levine, and Aviral Kumar. Digirl: Training in-the-wild device-control agents with autonomous reinforcement learning, Advances in Neural Information Processing Systems, 37:12461-12444 2024.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of artificial intelligence research, 47: 253-279, 2013.
Tommaso Bendinelli, Artur Dox, and Christian Holz. Exploring 1lm agents for cleaning tabular machine learning datasets. arXiv preprini urXiv:2503.06665 2025.
Jun Shern Chan, Neil Chowdhury, Oliver Jaffe, James Aung, Dane Sherburn, Evan Mays, Giulio Starace, Kevin Liu, Leon Maksin, Tejal Patwardhan, et al. Mle-bench: Evaluating machine learning agents on machine learning engineering. arXiv preprint arXiv:2410.07095, 2024.
Harrison Chase. LangChain, October 2022. URL https: /ttps:://ghthtbb/ com www..//cco//aaniinnaaaaaa langchain.
Sanjiban Choudhury. Process reward models for Ilm agents: Practical framework and directions. arXiv preprint arXiv:2502.10322 10325, 2025.
Paul Christiano, Jan Leike, Tom B. Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. 2023. URL https: wrw./rrg/ssss... 1706. 03741.
Quy-Anh Dang and Chris Ngo. Reinforcement learning for reasoning in small Ilms: What works and what doesn't. arXiv preprint arXiv:2503.16219 2025.
Federico Errica, Giuseppe Siracusano, Davide Sanvito, and Roberto Bifulco. What did do wrong? quantifying Ilms' sensitivity and consistency to prompt engineering. arXiv preprint arXiv:2406.12:: 12334, 2024.
Antoine Grosnit, Alexandre Maraval, James Doran, Giuseppe Paolo, Albert Thomas, Refinath Shahul Hameed Nabeezath Beevi, Jonas Gonzalez, Khyati Khandelwal, Ignacio Iacobacci. Abdelhakim Benechehab, et al. Large language models orchestrating structured reasoning achieve kaggle grandmaster level. arXiv preprint arXiv:2411.0355 2024.
Yang Gu, Hengyu You, Jian Cao, Muran Yu, Haoran Fan, and Shiyou Qian. Large language models for constructing and optimizing machine learning workflows: A survey, arXiv preprint arXiv:2411.10477 2024.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in lims via reinforcement learning. arXiv preprint arXiv:2501.12945. 2025.
Danijar Hafner, James Davidson, and Vincent Vanhoucke. Tensorflow agents: Efficient batched reinforcement learning in tensorflow. arXiv preprint arXiv:1709.02877 2017.
Matthew Hoffman, Bobak Shahriari. John Aslanides, Gabriel Barth-Maron, Nikola Momchev,
Danila Sinopalnikov, Piotr Stanczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, et al. Acme: research framework for distributed reinforcement learning. arXiv preprint arXiv:2006.0090 00999, 2020.
Jie Huang, Xinyun Chen, Swaroop Mishra, Huaixiu Steven Zheng, Adams Wei Yu, Xinying Song, and Denny Zhou. Large language models cannot self-correct reasoning yet. arXiv preprint 10.00:2310.01795 2023a.
10

---

Published as a conference paper at ICLR 2026
Qian Huang, Jian Vora, Percy Liang, and Jure Leskovec. Mlagentbench: Evaluating language agents on machine learning experimentation. arXiv preprint arXiv:2310.03302,, 2023b.
Zhengyao Jiang, Dominik Schmidt, Dhruv Srikanth, Dixing Xu, Ian Kaplan, Deniss Jacenko, and Yuxiang Wu. Aide: Ai-driven exploration in the space of code. arXiv preprint 2025.
Sham M Kakade. A natural policy gradient. Advances in neural information processing processinnnpoocssin systems, 14, 2001.
Sham Machandranath Kakade. On the sample complexity of reinforcement learning. University of London, University College London (United Kingdom), 2003.
Harrison Lee, Samrat Phatale, Hassan Mansoor, Kellie Ren Lu, Thomas Mesnard, Johan Ferret, Colton Bishop, Ethan Hall, Victor Carbune, and Abhinav Rastogi. Rlaif: Scaling reinforcement learning from human feedback with ai feedback. 2023.
Ziming Li, Qianbo Zang, David Ma, Jiawei Guo, Tuney Zheng, Minghao Liu, Xinyao Niu, Yue Wang, Jian Yang, Jiaheng Liu, et al. Autokaggle; A multi-agent framework for autonomous data science competitions. arXiv preprint arXiv:2410.206 20422444, 2024.
Eric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gonzalez, Michael Jordan, and Ion Stoica. Rllib: Abstractions for distributed reinforcement learning. International conference on machine learning, pp. 3053-3062. PMLR, 2018.
Zujie Liang, Feng Wei, Wujiang Xu, Lin Chen, Yuxi Qian, and Xinhui Wu. I-mets: Enhancing agentic automl via introspective monte carlo tree search. arXiv preprint 2rXiv:2502.14665 2025.
Siyi Liu, Chen Gao, and Yong Li. Large language model agent for Syper-parameter optimization. arXiv preprint arXiv:2402.0/86 0/81., 2024.
Xiao Liu, Hao Yu, Hanchen Zhang, Yifan Xu, Xuanyu Lei, Hanyu Lai, Yu Gu, Hangliang Ding, Kaiwen Men, Kejuan Yang, Shudan Zhang and Xiang Deng. Aohan Zeng, Zhengxiao Du Chenhui Zhang, Sheng Shen, Tianjun Zhang, Yu Su, Huan Sun, Minlie Huang, Yuxiao Dong, and Jie Tang. Agentbench: Evaluating llms as agents, 2023. URL https: //at:///axii...//i http//aawai.ggaaaaaaeeeeeee 03688.
/ttp/aa/////eeaacoee/.e//aeeaae/ abs/////////// /238.
Dakota Mahan, Duy Van Phung, Rafael Rafailov, Chase Blagden, Nathan Lile, Louis Castricato, Jan-Philipp Fränken, Chelsea Finn, and Alon Albalak. Generative reward models. arXiv preprint arXiv:2410.1283,, 2024.
Karthik Narasimhan, Tejas Kulkarni, and Regina Barzilay. Language understanding for text-based games using deep reinforcement learning. arXiv preprint arXiv:1106. 08941, 2015.
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730- 27744, 2022.
Jiayi Pan, Yichi Zhang, Nicholas Tomlin, Yifei Zhou, Sergey Levine, and Alane Suhr. Autonomous evaluation and refinement of digital agents. arXiv preprint arXiv:2404.0m 06474, 2024.
Tidor-Vlat Pricope. Hardml: A benchmark for evaluating data science and machine learning knowledge and reasoning in ai. arXiv preprint arXiv:2501.1566 2025.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley& & Sons, 2014.
Rafael Rafailov, Archit Sharma, Erie Mitchell, Christopher D) Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems,
2023.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-1897. PMLR, 2015.
11

---

Published as a conference paper at ICLR 2026
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: A flexible and efficient rlhf framework. arXiv preprint arXiv: 1095,, 2024.
Significant Gravitas. AutoGPT. URL htt.l//ps/ggghaooc.cc/////aaaaaa AutoGPT.
Shivam Singhal, Cassidy Laidlaw, and Anca Dragan. Reliability-aware preference learning for Ilm reward models.
Charlie Snell, Ily Kostrikov, Yi Su, Mengjiao Yang, and Sergey Levine. Offline rl for natural language generation with implicit language q learning. arXiv preprint arXiv:2206. / 1871, 2022a.
Charlie Snell, Mengjiao Yang, Justin Fu, Su, and Sergey Levine. Context-aware language modeling for goal-oriented dialogue systems. arXiv preprint arXiv:2204. 10198, 2022b.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can more effective than scaling model parameters. arXiv preprint arXiv:2408.033// 2024.
Guijin Son, Hyunwoo Ko, Hoyoung Lee, Yewon Kim, and Seunghyeok Hong. Lim-as-a-judge & reward model: What they can and cannot do. arXiv preprint arXiv:2409.11253 2024.
Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Xiangru Tang, Yuliang Liu, Zefan Cai, Yanjun Shao, Junjie Lu, Yichi Zhang, 7exuan Deng, Helan Hu, Kaikai An, Ruijun Huang, et MI-bench: Evaluating large language models and agents for machine learning tasks repository-level code. arXiv preprint wxXiv2231.009833 2023.
Sebastian B Thrun. Efficient exploration reinforcement learning. Carnegie Mellon University, 1992.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
David Venuto, Sami Nur Islam, Martin Klissarov, Doina Precup. Sherry Yang, and Ankit Anand. Code reward: Empowering reinforcement learning with vims. arXiv preprint 2xXx:2442.00474. 2024.
Siddharth Verma, Justin Fu, Mengjiao Yang, and Sergey Levine. Chai: A chatbot ai for task-oriented dialogue with offline reinforcement learning. arXiv preprint wwXiv:2204.08-22 2022.
Xingyao Wang, Boxuan Li, Yufan Song, Frank F Xu, Xiangru Tang, Mingchen Zhuge Jiayi Pan, Yueqi Song, Bowen Li, Jaskirat Singh, et al. Openhands: An open platform for software develop- generalist agents. In The Thirteenth International Conference Learning Representations, 2024.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8(3):279-29) 1992.
Yuxiang Wei, Olivier Duchenne, Jade Copet, Quentin Carbonneaux, Lingming Zhang, Daniel Fried, Gabriel Synnaeve, Rishabh Singh, and Sida I Wang. Swe-rl: Advancing 11m reasoning via reinforcement learning on open software evolution. arXiv preprint arXiv:2502.18449, 2025.
Ronald J Williams. Simple statistical crudient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229-256, 1992.
12

---

Published as a conference paper at ICLR 2026
Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, and Yiming Yang. Inference scaling laws: An empirical analysis of canpute-optimal inference for problem-solvingg with language models. arXiv preprint Ttt1:220080002244, 2024.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35:15476-15488, 2022.
Lunjun Zhang, Arian Hosseini, Hritik Bansal, Mehran Kazemi, Aviral Kumar, and Rishabh Agarwal. Generative verifiers: Reward modeling as next-token prediction. arXiv preprint arXiv:2408. / 52240, 2024.
Shujian Zhang, Chengyue Gong, Lemeng Wu, Xingchao Liu, and Mingyuan Zhou. Automl-gpt: automatic machine learning with gpt (2023). URL https://arxiv. org/abs/2300.022499
Yunxiang Zhang, Muhammad Khalifa, Shitanshu Bhushan, Grant D Murphy, Lajanugen Logeswaran, Jackyeom Kim, Moontae Lee, Honglak Lee, and Lu Wang. Mire-bench: Can language agents solve machine learning research challenges? arXiv preprint 2500::2504099002, 2025.
Shuyan Zhou, Frank F. Xu, Hao Zhu, Xuhui Zhou, Robert Lo, Abishek Sridhar, Xianyi Cheng, Tianyue Ou, Yonatan Bisk, Daniel Fried and Uri Alon, and Graham Neubig. Webarena: realistic abss
web environment for building autonomous agents, 2024. URL 2307. 13854.
Daniel M. Ziegler, Nisan Stienmon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences, 2020. URL https::///// /arxiv. org/ abs/ /1999, 08593.
13

---

Published as a conference paper at ICLR 2026
Appendix
In this appendix, we provide additional details on the method (Appendix A), additional details on the experimental setups (Appendix B), prompts to LLMs (Appendix C), and additional experimental results (Appendix D).
A
ADDITIONAL METHOD DETAILS
A.1 ALGORITHM FOR DURATION-AWAL GRADIENT UPDATES AND ENVIRONMENT INSTRUMENTATION
Algorithm 1: Policy Gradient with Duration-Awar Gradient Updates and Environment Instru-
mentation
Input: An LM agent тa with policy parameters 0, learning rate 7, batch size B, sampling multiplier m, total training iterations N, code execution environment P. another copy of the LM env-inst, a dataset containing task descriptions D, an empty buffer for previous solutions Preev
for teertiin 1 to N do prompts S~>DU prev
Sample
Sample solution for each prompt from the policy a ~ To(-1s)
Perform environment instrumentation and execute solution s' Pls,, env-inst(a)) / executions complete, each of which takes At and emits a reward R[s,, env-inst (a))
Wait until
Compute duration weighted policy gradient: VoJ(mo) E At V (o(al)) ((s
Update policy parameters: 0 + 0 (
Update previous solutions: Drrev +- self-improve prompt given a and s'
A.2 THEORETICAL JUSTIFICATION FOR DURATION-AWARY GRADIENT
In Section 1. we provided intuitive math why deration--ewwiighhed gradient can counter the frequency bias. Below, we provide the proof of the unbiasedness of the policy gradient estimate with duration-aware gradient.
Let the standard reinforcement learning objective be the expected cumulative reward, J(0) = Ey [R(r)], where R(T) Emo rt is the return of trajectory 5 sampled under policy πg. The Policy Gradient Theorem states that the gradient of this objective is:
VoJ(O) = Berd""aatto()) (o)t((s)
Vo log (a)ss)) (sm
where d= (s) is the stationary state distribution induced by mo and A (s, a) a valid estimate the advantage function. Let p(s. a) d= (F)))TT(( (s)TT(a(ss) denote the true state-action visitation distribution under the policy.
actions, the learning agent does not receive samples drawn from p(s, a). Instead, it receives a stream of completed transitions where actions with shorter duration At(s, a) appear more frequently over fixed time period. The probability distribution of these observed samples, which we denote Pobs (8, a), is therefore inversely proportional to the action duration:
In the described distributed setting with asynchronous,
p(s,a)
Pobs (8, a) At(s,a)
Let C' (B(e.o)<p
be the normalization constant. Then the observation distribution is
Atts,a)
explicitly Pobs (s, a) A naive policy gradient estimate using samples from this biased
p(s.a)
observation distribution would not yield the correct gradient direction.
To correct for this sampling bias, we employ importance sampling. We aim compute expectation over the true distribution p(s, a) while using samples from our observation distribution Pora a). The required importance weight w(s, a) is the ratio of the target to the proposal distribution:
p(s,a)
p(s,a)
w(s a) Pona a)
C. (t(s,a)
p(s.o)
At(3.e)
14

---

Published as a conference paper at ICLR 2026
By weighting each sample from (s a) by w(s, a), recover unbiased estimate of the true policy gradient:
V VgJ(() (+) E (s,a)~p
Vo log
E (x.a)~Pata
w(s,a) a) Vologmo(aas))
C. Ai(ssa)) V (o)
The expectation term in the final line is precisely the quantity estimated by the duration-aware gradient update rule. The derivation shows that this estimator is proportional to the true policy gradient (0). The proportionality constant C' is positive and does not depend on 0, meaning an optimization step in this direction is valid ascent direction for the policy objective J(0), with the constant being absorbed by the learning rate. This confirms that the duration- -warre policy gradient provides unbiased estimate of the policy gradient's direction.
15

---

Published as a conference paper at ICLR 2026
B ADDITIONAL EXPERIMENTAL DETAILS
B.1 HYPERPARAMETERS
Table 3: Hyperparameters for RL training of MLE agent.
Hyperparameter
Value
nax_prompt_length
1024
max_response_lent
1024
train baiin_atch__ae size
128
lotal_epochs nnodes
100
1
n _eppu__er_nodeeee
8
actor poor_model_tyye
Qwen2.5-3B-Ing Instructt
actor_enabi gradient cerieppiicchck
True
128
ecto__po__micro_aacchhoo
8
actor_grate clip
1.0
actor_clip_rane
0,2
actor_entropy_coem
coeff
0.001
actor_ppooepochs
100
actor actor_learning_.na
le-5
refcrence_log__rrbbmiinnn battrr___btc_sttztt__t
8
rollout_name
vllm
temperature
0.7
top_k k
-1
top_p
promptJenngh
1024
response_length
1024
dtype
bfloat16
apu_memory_uulizatton
0.2
ignore_eos
False
enforce_eagee _aager
True
tensor cirtorrmoel__peee _ar_em___rralllllllceceeppe size
8
log l..proobmmcrre batchh size
8
cntticoppimmizr
critic_leamnnggrrat
le-5
critic_ critic_mndell_tye
Bren2.5-3B-Instrue
True
critic_ppo_minn sttpo_mmni_baatch_sizeeee
128
Balll_ppp__m__ii______aa___eee micro batch size
8
critic forward micro batch size
8
critic lll___pppppppppppeeeeee_eee
100
critic lelg__rrddiiii
1.0
critic_cli ciriiccliprange_vve value
0.5
ppo_gamma ppo_lam adv ddv_estimator ppo _o_c
1.0
1.0
gae
0.001
B.2 COMPUTATIONAL RESOURCES
We train the Qwen2.5-3B. --Istruct model until convergence. We train each task using 8 A100- 40GiB NVIDIA GPUs usinv the hyperparameters above. The amount of time until convergence depends the task, which varies from to 3 days.
16

---

Published as a conference paper at ICLR 2026
C EXAMPLE PROMPTS
C.I PROMPTS FOR SOLVING KAGGLE CHILLLNNGE
Example Prompt for random-actss-o--pizza
<lim_ sttsst>ts>sssem You are given a machine learning task. You must solve the task by train- ing a model and rumning the model on the test set to produce submission file.<lim_ent endl> <lim_ stn__ttrti>user First outline your proposed solution natural language (3-5 sentences), fol- lowed by single markdown code block (wrapped in "*). Note:
Datasets train.json and test.json are available in /wworkdir/randoo-actt--fff
pizza//prepared/pubiicc..
You MUST produce a submission file "/submission.csv' by running your model on the test split.
Your task is to predict the probability that request resulted in a free pizza.
Submissions are evaluated on the ROC curve between the predicted probability and the observed outcomes.
Your can use pre-installed packages such as: 'statsmodels', 'pandas', 'timm', 'bayesian- optimization', 'scikit-learn' 'xgboost', 'numpy', 'torch', 'torchvision', 'lightGBM', 'torch-geometric..
You can't t access the internet so don't I use any pre-trained models need downloading.
//sbmissionncsst should have the following format: quest_ hea__iddrequestrrreeeeee caatteee_reevveed_iiieeee 13 i8.44,81440 c3_ Imfqi000 etc
re-
Data snippet: ->
"giver_username_tt in_wn" "N/A",
"request_id": "t3_ "t3_law5zf",
"request_tedi text_eeeeeeee_eeeetteee edit aware": "Basically had unexpected expenses this month out money and desperate for pizza. Have Tera account with level 48 Beserker and the account has founder status. not much but only thing have right now that can live without. Eating is much higher my priority list right than playing Tera. If you don't ton't want the account I will be happy to pay it forward to someone this friday when get my paycheck.",
"request_titte _tttlsstttttt "[Request] Don't have much but willing to trade.".
"requester_account, account accoont_age__aaaaa_____ctaeea__e in atygussttr 165. (65.9420949074074 "requester_daye dayss_sicceee____aeeaaaaaaaaeea_
since first lrrss__pst__oge on can_rrop__at at request": 0.0.
httpare__ummebroo_oommmeeeeee attruuusttrrrrt" 13,
"requester_numberr Ltapeeettttnnnooopp__e_ee
_t__rro__a_rrreuus""" 0,
1,
"requester_number" of _l__pooss__ooo on raop_at_paaaeaaaeeaaaaaaaae at_rrttustt:: 0,
6,
"requssee__sbrrdddis_____ee
"TeraOnline".
"Torchlight".
"funny",
"pics",
"todayilearned". "windowsphone") ].
Dicquester_upvotes_miiiee setroote_mmiuss_eee catesssmmuuussuosoeeeeeeeee @oownoote__t_reqeeeee 168,
"requester_uupooes__lue fatlllusdowvooss___e "atprootts_rrrreues"""""" 240,
Trequester_usernamee "VinginityCollectoor.
"unix timestamp_po requestt" 1364094882.0.
"unix firx_timettamp_oft ree_e_ss__sss__tee" 1364091222.0...
<lim endl> <lim startbassistant
17

---

Published as a conference paper at ICLR 2026
C.2 PROMPTS FOR EEREVVENTV
Example Prompt for Self- Self-Improvement
<lim_start startlsssstem
You are given machine learning task. You must solve the task by training model and running the model on the to produce submission file.
<lim him_endl>
<lim_sttare startbbuser
You have implemented previous solution. Revise the solution to improve the performance on the test set. First outline your proposed solution in natural language (3-5 sentences), followed by single markdown code block (wrapped in "") which implements this solution. If you reuse parts of the example code, include those sections again in your final solution. Previous solution: code)" plan_code)"
"' prrevious plan_
<lim siin__end>>__
<lim him_sttrtbassistanttt
C.3
PROMPT FOR ENVIRONMENT INSTRUMENTATION
Environment Instrumentation
Please insert print statements the given python script. The print statements are supposed to reflect the progress of executing script that solves Kaggle challenge machine learning benchmark. These print statements will be used to debug the python script so it needs to capture the progress of execution.
Print Statements:
print( iimpooiimote packages")
print("loaded data")
print("defined model")
print("training loss:")
print("trainec model"))
print( tresingt loss:")
print(" prreddpddd test labels")
Requirements:
Only insert print statement AFTER operation is actually performed (e.g., data have actually been loaded).
Insert print statements for "training loss:" and "testing loss:" applicable (i.e., the code actually computes training testing losses).
Output the entire python script after inserting print statements in single markdown code block (wrapped "").
Do modify the original python code, other than inserting print statements.
Now please insert print statements for this python script: (code)
18

---

Published as a conference paper at ICLR 2026
D ADDITIONAL RESULTS
D.1 AVERAGE SCORES ACHIEVED DURING RL TRAINING
Average score (across 128 samples)
-1
-2
-4
-
-6
rand ti-ppeaazze
tearstr sageeg--- las-est aoutted -oany-sceriing
-7
tueeh-ee bert---eetiitt--ttrrtte
-8
eather iotettttttttiio
-9
steps
10
20
30
40
70
100
RL Training Steps
Figure 10: Average scores achieved during RL training different tasks. Scores are -10 for invalid solutions and actual score from grader in MLEBench (Chan et al., 2024) if scores are valid. RL consistently improves average scores across tasks and across seeds per task.
D.2 FULL RESULTS FOR ABLATING SELF--MPROVEMER PROMPT
Table 4: Improve Previous Solution leads to further improvement than only solving the task from scratch on 10 out of the 12 tasks.
Tasks
From Scratch
Improve Previous
deeerinngiialllssinscccccoooommn (t)
0.898 +/- 0.003
0.895 +/- 0.001
sarningggggnyylllaaauummdddeeees (f)
0.729 +/- 0.004
0.746 +/- 0.002
random--acts-of-pizza1 (1)
0.643 +/- 0.004
0.663 +/- 0.011
(weetsseentimmntleettaccioo (1)
0.593 0.000
0.596 +/- 0.002
sabularrplaggooudddeereee--aa 2022 (t)
0.902 +/- 0.000
0.913 +/- 0.000
abbllarrplyyroomdddseeeseseeeeed (1)
0.950 +.000001
0.951 +/- 0.000
sxppacll-pirraee---plnn-mmattttt (t)
0.517+/-00
0.527 +/- 0.003
plnttppatholog-22220-gewr (t)
0.949 +/-0.017
0.970 +/- 0.004
leaf-classification ())
0.469 +/- 0.244
0.124 +/- 0.000
nomad2018-pp 8--rrddcct-traneen weapeett-traaaaaaattooooucccttte (4)
0.060 +/- 0.000
0.059 +/- 0,000
(4)
0.448 +/- 0.000
0,404 +/- 0.011
imsys-chatbot-arenaa (L)
1.098 +/- 0.003
1.081 +/- 0.002
19

---

Published as a conference paper at ICLR 2026
D.3 EFFECT OF SAMPLE SIZE
MLEBench limits each run to 500 and 5000 solutions in the AIDE tree for the 24hr and 100hr run, respectively. We tried relaxing this limit to allow for a max 1000 solutions for the 24hr run but did not observe better performance, reporting. Below is the statistics of the number of solutions for few example tasks from running AIDE for 24hrs.
used the same setting for AIDE in MLEBench when
Task
Mean
Std
Min
Max
445.00
373.08
62
1000
clafflaassifiiation 277.00
377.89
12
1000
nomad2018 1omaa2201-preddiitttraaaaap Random-actt--f-pizza
681.00
438.46
6
1000
317.00
306.35
78
1000
Table 5: Statistics of the number of solutions for few example tasks from running AIDE for 24hrs.
As seen from Table 5, the max performance of Claude+AIDE is generally aggregated over a few hundred of solutions at most. Hence we did not want to exceed this number when aggregating over the number of samples. The choice of 128 was also partly determined by RL training, since this is generally when the base Qwen model could sample valid submission (for RL to make progress). Below, we report the effect of sample size after RL training for a subset of the tasks in Table 6
number samples
detecting-insun insults learning-agenn
random-acts-of-ppzzz
0.862
0.679
0.621
0.893
0.75
0.663
128
0.895
0.746
0.663
256
0.895
0.752
0.661
Table Best performance under different number of samples after RL training for subset of the tasks.
We found that the max performance is similar for sampling sizes of 64, 128, and 256. When the sampling size is reduced to 8, there is significant performance drop.
D.4 ALTERNATIVES RL
We have indeed tried many different finetuning approaches to adapt the small model to particular MLE task, including SFT, DPO, RWR (reward weighted regression), and STaR. None of these offline method worked well running online RI.. Below is table of running STaR (sampling bunch of trajectories, filter high-performing trajectories to finetune the model on) and how it compared to the original Qwen model:
Base Qwen
Filter for top 50%
Filter for top 90%
0.60 +/- 0.08
0.59 +/- 0.12
0.62 +/- 0.14
Table 7: Performance random-acts-of-pizza running STaR Zelikman et al. (2022) algorithm, which was uneffective for MLEE
D.5 ADDITIONAL QUALITATIVE EXAMPLES
D.5.1 CODE EXAMPLE BEFORE AND AFTER ENVIRONMENT INSTRUMENTATION
20

---

Published as a conference paper at ICLR 2026
Import pandos pd
from sklearn. wnatuure_etttattioonooo inport TridfVectoriaer
Erem skleacn.moder model_selectton import tsain_test__spli
frem sklearn. ensemble import Kandmmocesttlasssifee from sklearn.metries import roq_auc__core
import nunpy sa np
Toad and preproceas data train_date - pddreedd j acta of ptaps//ppepprre//ppuullcooo json*)
10 test_date pd. ("/oo///r///ddrraaaadaa -at-aaadnnn--o-oo--------o http::/..pp.//ppppps/////ll/ccce jaon")
11
12 Preprccess
13 Wiid__eecoorie TfidfVectorizer (stop w
max_[eeturea=1000
X.train_tfidf weidf_wectozizer.. ftt. "atta""equest_tees
15 X_test_tTior
www.aaeeeeettttmmmmmmm...mmmma (traattttttttttttttttt ["repusst_ext__edtt_t
16
17 Train kandmeereesclassifie:
18 =f_claasifiar = kmdomsoraattcaaaciiieee htw.://aaat..iioootttttttrrooeoe sanden_state==2)
20 __classiiier. www.ttttt....... [X_irrii_ttiddr train_data, "tegpeseeerr
22 Predict probabilities for the test set
23 y-rore_proba S__classffeer.predicttto (X [Ltest_tfid]] I:.
25 Prepare sabiisaion file
submission
27 ("request_id", Lest_data("=eedee "tss_datt("=eques___e""] "requesterr yrobed_probal
29
submi ssion. to cov "oov""./aubmaaeeee wwwp///aaiccnc.ommm index=Falsa)
Figure 11: Example code before environment instrumentation for solving the random-acts Kaggle task.
Import pandas
Zrom skleazn. www.eeeettatiionnncooo Lext inport TridfVectorizer
Erem sklearn.modes stapssnnodd___seeeeeee
import train_test_splle
from sklearn. ensemble import kundom/orastCaaaaifio.
sklearn.ensemble
skleacn.metries import rcs_auc_score
import mumpy za
print ("imported packages")
10 Load and preprocess data
11 train_data = read_
teat_data- pol.readlsoonn"///rrrr////ooaa http:////annnooooooeera/a/////rr
13
print loaadd data")
15
16 Preprocess text data
17 tfidf_veetorizer "Iiidivactorizer (trprrrrrrrrrrooooooodoooooiaiia mox_Faaturos=1000)
18 X_tran__ffaf tfidf_ t"ldf_ vectorizeer.r.. zer, fit traanaarrmttaail (train_datt data aa cecppest__ttt__dtttt_.
X-test_ffidr LLidr_rectoriser. (ta. ["request__tttttt_____
20
Train a RandomForestClassittee =fllassiffer RandonForesttClessseee htw.::/aeeee.eii....eeetttteoeor sandon_state-(2)
22
25
print ("dafined madel"]
25
=[elassifier. Mittriin Lrain Leiiiaattt/rrpppserreeeeeeeeaas. 27
28
print ("izained model")
30
Predict pichabilitties for the test set
y_pred_prooa rf_elaast farrrffrrppaar ct_proba (X_test_tffdd)) [t,
32
print ("predicted lebels")
34
35 Prepare submission file
36
submission aataraare
37
("requeat_id";
troppestor_renntvve_ppe prored_prohal
39
submis cilmission.do csv. //w__ccw../uuummeeiioooooo Index=Palse)
Figure 12: Example code after environment instrumentation for solving the random-acce continaattttt-ppppp
Kaggle task. All print statements generated by environment instrumentation.
21

---

Published as a conference paper at ICLR 2026
D.5.2 EXAMPLE TRAINING
SOLUTIONS FOUND BY THE AGENT AFTER RL
areplan: By proposed solution involves training a machine learning model using a combination. request titles descriptions, along user-specific featurea account number of posts, etc. will proproceas the text data using TF-ThF vectorisation and coacalenate it with the aumerical Seatuzes, The model will be a canden fores: classifies handle the binary classification task of predicting whether requeat will reault in free plzza, After training the model on the training set, I will
fine its hyperparameters using Beyesien optimization,
to predict
the probability of free pizza for each request in the test set.
# Import necessary libraries
import pandas as pd
Erem dllearn.ensemble import Kandmmoozsstcaasssiiie
from sklearn. feature_ pettureeeexxtcteee rxtnaccionn import Tridffectorizer
sklearn.model www.emmdeeeeeittii...... import GridSearchCV import mumpy 25
10 Load and preproceas data
11 train_dE pd. json /wo.kdir/tanddm acts of pttza//pppperee//oi//t/aaaa ison')
test_af pd.read_ison /
13
seature engineering
15 tfidf
max mmx_featture=10001
16 (traain_ request reeueeet_exx__eeeeeeeeeeee ter_ett__et__ell__aareerrr
17 train_tfidf
bettt__eddt_wwasee_ee (f
18 test_tfidE SiidE.transfomlteett.cc. http::ddde_eeeeeeccccecccoooooo.
20 Combine features with user-specific festures
train_festuree mp.hsteck (frraaaaaaallllloossssssssssooe
22
23
24
25
27
29
valess,
30
test_features
32
34
35
36
"ecuester__uuber__f_pott.
37
38
"=oquester_uuvotes_nue htt.://grroooo.ooooodeeeeeeeeees htt_aauo__eoeeeee____eeeeeee
39
http::/e....poooooooooooootttaoe vavuues,
test_ifidf,
41
42
model
4 zfc RardoniFrrest ineeeaiissee :acdon_etate=12))
44 parem_grid ('max depth'' ap. ap.arange 15), "min_samples_iedf" mp.azange 191)
gnar_ Gridscarchcv paren_grid, ev=3, sborings/ros_uuee
46 grad__earch.dfrr (train_ features, train_df(' p
88 # Predict lities en test set
49 teat_proba prrii_eeacch.predir_..o [teat_Ceaturea) [t, 1]
30
31 Prepare submission file submi submission_af nataFrame request_id'; test_df ('rquaaat_id") 'roquester_ree recei piza'' test test_probal! butmission_d..toccrm w/w.mmissiioncc..// index=False)
Figure 13:
An example of high-performance (score 0.66) high cost (115 seconds) solution for cindom-acts-cnn of-iizza achieved after RI. training.
22

---

Published as a conference paper at ICLR 2026
preplan: The proposed solution invelves preprocessing the train and cest datasets to extract on extracted feasures, then predisting acores relevant festures, training for the test set. For this task, will use transformer-baaed model I from the 'timm" library to process the text data, and a 1ightGBM claasifier roz the predict ion Lask, The from the train data and then
model will he traired the features evaluated on the test aet to generate the aubmlasion file,
import pandas as pd import numpy np from sklearn,moce maaern.oddel_eeeeteee
import train_test_split
frem 1ightghm import LGBKClaseifier
xgboost Import xcsclaast
from skleazn. wealuureextzattttnn
import TIid[eectorizer
tgém.notebook import todm.
10
Toad
12 train_d[
agency- lab-lalb-lle automated soarg
/tzain.cev") pd.ppedddddeee (ed.rreddccoo///oooddr//lllaa aaaeaaa--nn-----eeeeeeee /at
:est.csv")
15 Excract Text Features
16 riitr_vectorierr TfidfVectorizer()
17 train_tfidi
18 test_tfidE
text ")
20 Train Model
Lnaxclasaifier wtw.raaierrrraaiiiirrrrrol.. max_depth=10)
22
madel.fit (train_t*idf, (rain_
24
Predict Test Set test_preds = model.predict (test_tfide)
25
27 Save Sutmission
DataFrame
29 *esaay_ id"s gest_dt.essay..
30
scoze": Lest_preds
31
32 submi cobmission.to.o ("sumiiiiiioo cav", index=False)
Figure 14: An example of perforrrrme (score 0.73) high cost (281 seconds) solution using gradient boosting for Learning-agency- wub-auttmattedeeeeyyaaaaaaaaaaee achieved after RL training.
23


Reviews:
[{"review_id": "CIGUZadDzL", "reviewer": "Reviewer_kZcG", "summary": "This paper analyzes the best way to improve performance on a MLE task with a well-defined metric and grader, specifically focusing on the benefits of prompting/scaffolding versus incremental RL. The authors explore duration-aware gradient updates and environment instrumentation / partial credit and ultimately find that a small model with RL performs best in this setting.", "questions": "How do these results generalize to settings beyond this 12-task set?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors focus on a realistic problem (machine learning engineering) and ask a very timely and relevant question (whether it's better to prompt/scaffold vs train with RL to improve performance on a task). This is a question many application businesses are asking themselves right now. The authors prove their points with practical and sound experiments that are clearly laid out. They test a variety of models, scaffolds and settings and lay out results clearly.", "weaknesses": "The authors focus on 12 tasks; it would be interesting to increase size (and specifically, diversity and generality) of the task set.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C6hbmfnOTs", "reviewer": "Reviewer_DyMM", "summary": "This is a great paper. It does RL to solve machine learning engineering tasks. From what I understood, the problems with those tasks is that when formulated as MDPs, credit assignment is difficult because the only non-trivial reward is given at the end of a trajectory when the final engineering task is solved or not. To get intermediate rewards and mitigate this problem, authors propose environment instrumentation (adding prints). They also propose duration aware gradient to favor high-cost but high-reward actions.", "questions": "Could you try GAE estimatation a.k.a elligibility traces in your RL experiments for even better credit assignment? \nCould you explain why asynchronicity is needed? And if not needed, would duration aware gradients still be needed?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "This paper's ideas are great and could make it a seminal paper for anyone interesting in RL for machine learning engineering. \nThe experiments are amazing with comprehensive ablations and multiple seeds of RL training which is rare in nowadays research.\nThe problem inherent to RL for MLE are well presented.", "weaknesses": "Here is the thing, the paper is clear as is and has very few objective weaknesses. While the propose duration aware gradients and environment instrumentation are objectively (supported by strong experiments) good contributions, I also believe that objectively the experiments do not allow to determine if RL for MLE is a real advantage. Correct me I am wrong, and I am not an expert in LLMs, but it is semms that authors only use RL for the qwen model and compare the latter to e.g. claude. To really show RL benefits for MLE, wouldn't it be better to show for a set of K models the performance differences between k_1 k_2 k_3... and k_1+RL, k_2+RL ... ?\n\n\nThen, from an RL perspective, it seems to me that both duration awareness and environment instrumentation could be presented, not algorithm tricks but as MDP formulation changes: duration awareness could be added to the reward function, and prints could be added to the states and transitions? Also the asynchronous setting is not clearly motivated. It would better to focus on the fully synchronous setting for now.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hEmy5jCrDp", "reviewer": "Reviewer_puSq", "summary": "This paper explores how to best utilize compute in machine learning engineering (MLE) tasks. The paper proposes to adapt a small LM through reinforcement learning (RL), leveraging duration-aware gradient updates from the task’s verifier signal. Empirical results on 12 Kaggle-style MLE tasks show that an RL-adapted Qwen2.5-3B model outperforms a much larger static model (Claude 3.5 Sonnet) using the state-of-the-art agent scaffold (AIDE) by an average of 22%。", "questions": "1. Could you report performance with only duration-aware updates, and only partial-credit instrumentation, to isolate their contributions?\n2. Does the RL-adapted agent transfer to unseen Kaggle tasks without additional fine-tuning?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Technical contributions: Both duration-aware gradient updates and partial-credit instrumentation are conceptually clean, easy to implement, and empirically validated.\n2. Strong empirical evidence: Demonstrates small adaptive models outperforming much larger static ones, which is an impactful and counterintuitive result.", "weaknesses": "1. Lack of detailed ablations: The effect of partial-credit feedback is not isolated in the reported experiments.\n2. Generalization and robustness unclear. It is not evident whether the adapted model transfers well to unseen MLE tasks or remains stable under changing task distributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DUQUXG95mf", "reviewer": "Reviewer_jqr7", "summary": "This paper explores a RL framework for code generation that incorporates environment instrumentation to provide more informative rewards. By measuring both execution time and partial correctness of intermediate steps, the approach aims to make RL-based fine-tuning of code models more efficient and stable. The method is implemented on code generation benchmarks, and experimental results show that a small model trained with RL can outperform larger models or those using scaffolding, demonstrating the practical effectiveness of the proposed approach.", "questions": "Please check my comments in Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The idea of integrating execution time and partial step validity into the RL framework is innovative and interesting, providing a richer and more informative learning signal.\n2. The paper is well-organized and clearly written, with logical flow and easy readability.\n3. Experimental results convincingly show that a small model enhanced with RL can outperform larger models with scaffolding, validating the method’s efficiency and effectiveness.", "weaknesses": "1. The proposed Environment Instrumentation is conceptually interesting but appears to be an engineering optimization rather than a fundamentally new RL formulation. It would be helpful to explain why simpler signals, such as the error line number from runtime exceptions, could not serve as reward feedback.\n2. The experiments lack comparisons with other RL-based code generation methods, focusing only on commercial or scaffolding-based baselines. This limits the clarity of how much the proposed method improves over existing RL techniques.\n3. Minor issues: Line 245 contains a possible error, and Figure 4 seems mislabeled (the prompt content is missing and should likely correspond to Table 4).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)