Paper deep dive
Reward-rate Policy Gradient for Efficient Machine Learning Engineering Agents
Muhang Tian, Sherry Yang
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Traditional reinforcement learning (RL) techniques focus on maximizing expected cumulative reward, where each action assumes to take a constant unit of time. However, this assumption does not hold for agentic RL tasks such as machine learning engineering (MLE) agents, where actions involve data loading, feature engineering, and model training that take variable durations. Efficiency matters in modern agentic RL where actions are costly. To address this limitation, we adapt from continuous-time RL and Semi-Markov Decision Process (SMDP) formulation and propose Reward-rate Policy Gradient (RPG), where we focus on optimizing the reward rate -- the long-term reward per unit of time. RPG estimates the reward rate from off-policy samples, then charges each action for the time it consumes at that rate. We first conduct theoretical analysis in the bandit setting to establish that RPG approximates the optimal reward rate and empirically demonstrate it outperforms baselines while avoiding enumeration over the policy space, a known issue for an existing method. We then further apply RPG on a small language model (Qwen3.5-4B) with self-improvement loops and empirically show it obtains higher rewards within a fixed time budget than vanilla RL on MLE-Bench and NanoGPT, with a 19.2% and 85.7% margin, respectively. Our method provides a practical solution for optimizing performance under wait time considerations in modern agentic RL tasks, where actions interact with external environments and cost time.
Tags
Links
- Source: https://arxiv.org/abs/2609.36393v1
- Canonical: https://arxiv.org/abs/2609.36393v1
Trouble viewing inline? Open PDF directly →
Full Text
144,442 characters extracted from source content.
Expand or collapse full text
Reward-rate Policy Gradient for Efficient Machine Learning Engineering Agents Muhang Tian Sherry Yang Affiliation: New York University Email: mt4193@nyu.edu Abstract Traditional reinforcement learning (RL) techniques focus on maximizing expected cumulative reward, where each action assumes to take a constant unit of time. However, this assumption does not hold for agentic RL tasks such as machine learning engineering (MLE) agents, where actions involve data loading, feature engineering, and model training that take variable durations. Efficiency matters in modern agentic RL where actions are costly. To address this limitation, we adapt from continuous-time RL and Semi-Markov Decision Process (SMDP) formulation and propose Reward-rate Policy Gradient (RPG), where we focus on optimizing the reward rate — the long-term reward per unit of time. RPG estimates the reward rate from off-policy samples, then charges each action for the time it consumes at that rate. We first conduct theoretical analysis in the bandit setting to establish that RPG approximates the optimal reward rate and empirically demonstrate it outperforms baselines while avoiding enumeration over the policy space, a known issue for an existing method. We then further apply RPG on a small language model (Qwen3.5-4B) with self-improvement loops and empirically show it obtains higher rewards within a fixed time budget than vanilla RL on MLE-Bench and NanoGPT, with a 19.2% and 85.7% margin, respectively. Our method provides a practical solution for optimizing performance under wait time considerations in modern agentic RL tasks, where actions interact with external environments and cost time. †footnotetext: Website: https://muhang-tian.com/reward-rate/ Code: https://github.com/MuhangTian/Reward-rate-Policy-Gradient 1 Introduction RL has been used extensively to train language models (LMs) and agents, where the standard approach is to maximize expected cumulative reward (Christiano et al., 2017; Zhang et al., 2025a). Recently, there has been interest in MLE agents, where the goal is to automate the MLE pipeline to achieve best performance on held-out data (Chan et al., 2025; Yang et al., 2025; Nam et al., 2026). Current approaches mostly apply a strong LM in a search procedure: AIDE explores the space of candidate programs as a tree (Jiang et al., 2025), while MLE-STAR retrieves an initial solution with a search engine and then refines its targeted components (Nam et al., 2026). A separate line trains the model itself with RL on performance score (Yang et al., 2025). However, none of these methods perform gradient updates that fundamentally change agents’ behavior to learn to generate more efficient solutions. In MLE, each action takes variable time and resources, so it is important to consider samples’ costs for practicality — we want agents that are not only good, but also efficient. One solution is to use multi-objective RL and search for the Pareto frontier between performance and time (Roijers et al., 2013; Hayes et al., 2022), but computing the set of optimal solutions is often infeasible and compute-heavy (Hayes et al., 2022). These limitations motivate an objective that intrinsically aims for long-run efficiency. One natural formulation is the long-run reward per unit of time studied in Semi-Markov Decision Process (SMDP) (Sutton et al., 1999b; Das et al., 1999; György et al., 2007). Intuitively, the method uses RL to maximize relative rewards, calculated from total rewards minus the amount that it would have earned if time were spent optimally. However, performing this computation requires knowing the optimal reward rate apriori, which is often not the case in practice, giving a chicken-or-egg dilemma. Existing approaches mainly use two ways to tackle this. Directly optimizing the reward rate with SMDP policy gradient objective (Sutton et al., 1999a) requires estimating the rate of the current policy from on-policy rollouts, an operation that is still expensive when actions are costly. The continuous-time bandit approach (György et al., 2007) estimates the optimal reward rate from historical samples, but it requires an exhaustive search over the policy space, infeasible in modern agentic RL settings where action space is large. To that end, we propose Reward-rate Policy Gradient (RPG), which uses off-policy historical samples to estimate the reward rate under the greedy policy and maximizes for relative reward (Figure 1). RPG avoids the limitations of the two previous methods. It does not need time-costly on-policy rollouts for each gradient step, and does not require searching over the policy space. Additionally, RPG is simple to adopt in the current RL pipeline, since it only requires a reward-rate estimator and uses the relative reward as the learning signal. We first study RPG in the bandit setting and theoretically prove that, in expectation, if the greedy policy’s reward rate is estimated well, then RPG is equivalent to the Dinkelbach iteration (Dinkelbach, 1967; Schaible, 1976), and natural policy gradient (NPG) (Amari, 1998) is suited for reward-rate optimization due to preconditioning on Fisher information. Experiments on the bandit setting confirm that RPG obtains lower regret than the reward-rate upper-confidence-bound algorithm by György et al. (2007) without the need to search over the policy space. These findings motivate us to apply Proximal Policy Optimization (PPO), a first-order approximation of NPG methods (Kakade, 2001; Schulman et al., 2015), on MLE-Bench and NanoGPT. Compared against vanilla RL baselines on the same time budget, RPG shows a 19.2% average improvement in performance across 20 out of 22 MLE-Bench tasks and an 85.7% margin on NanoGPT. Our findings demonstrate that RPG is a practical solution for future agentic RL frameworks where actions are costly. Figure 1: Overview of RPG in MLE, where our goal is to optimize for reward per unit of time in agentic RL contexts. Our agent under policy πθ _θ is asked to self-improve upon its previous MLE scripts from buffer ℬB, and each action a is a proposed script that executes in variable time δ and admits a performance score r such as AUC. ρ ρ is estimated from past samples in buffer ℬB and captures the cost per unit of time. Relative reward r~ r is then used in PPO to perform RL training. Maximizing the relative reward optimizes for reward rate, giving higher reward per unit of time and more efficient RL agents. 2 Preliminaries 2.1 Semi-MDP Policy Gradient We introduce the notation and setup for the SMDP setting where each action takes a variable amount of time (Sutton et al., 1999b; Das et al., 1999). Let states be s∈s∈S and actions be a∈a∈A. At each step t, the policy π(at∣st)π(a_t s_t) is a conditional distribution on the actions. Taking the action ata_t gives reward rt+1:=r(st,at)r_t+1:=r(s_t,a_t) and time δt+1:=δ(st,at) _t+1:=δ(s_t,a_t), and the transition model is a joint conditional distribution p(st+1,δt+1,rt+1∣st,at)p(s_t+1, _t+1,r_t+1 s_t,a_t). Our objective is to maximize the long-term reward per unit of time: ρ(π)=limt→∞π[∑j=1trj]π[∑j=1tδj],π∗=argmaxπρ(π).ρ(π)= _t→∞ E_π [ _j=1^tr_j ]E_π [ _j=1^t _j ], π^*= _πρ(π). (1) For a parameterized policy πθ _θ, taking the gradient gives ∇θρ(πθ)=∑sdπθ(s)∑a∇θπθ(a∣s)Qπθ(s,a)πθ[δ], _θ\,ρ( _θ)= _sd _θ(s) _a _θ _θ(a s)Q _θ(s,a)E_ _θ[δ], (2) where Qπθ(s,a)=πθ[r−ρ(πθ)δ+Vπθ(s′)∣s,a]Q _θ(s,a)=E_ _θ [r-ρ( _θ)\,δ+V _θ(s ) s,a ] is the differential action-value function (see full derivation in Section A.1). Empirically, the division by πθ[δ]E_ _θ[δ] can be dropped since modern adaptive optimizers like Adam (Kingma and Ba, 2015) and gradient clipping reduce its impact on the gradients. Thus, the implementation is a straight-forward plug-in using the relative reward r~t r_t into existing RL algorithms r~t r_t :=rt−ρ^tδt,et=r~t+Vϕ(st+1)−Vϕ(st),A^t=∑l≥0λlet+l, :=r_t- ρ_t\, _t, e_t= r_t+V_φ(s_t+1)-V_φ(s_t), A_t= _l≥ 0λ^l\,e_t+l, (3) ℒ(θ) (θ) =t[minwt(θ)A^t,clip(wt(θ),1−ϵ,1+ϵ)A^t],wt(θ)=πθ(at∣st)πθold(at∣st), =E_t [ \w_t(θ) A_t,\;clip (w_t(θ),1-ε,1+ε ) A_t \ ], w_t(θ)= _θ(a_t s_t) _ _old(a_t s_t), (4) where ρ^t ρ_t is a Monte Carlo approximation of ρ(π)ρ(π) from on-policy rollouts. Intuitively, the relative reward r~t r_t is the gain obtained for the selected action when δt _t units of time are used at a cost of ρ(π)ρ(π), so actions are ranked by their relative performance per unit of time. 2.2 Continuous Time Bandit Reward-rate RL has also been studied in the continuous-time bandit setting (György et al., 2007). The agent observes context xtx_t drawn i.i.d. from a fixed distribution on a finite set X, and each action at∈a_t∈A yields a reward rt+1r_t+1 and time δt+1 _t+1, with r∈[rmin,rmax]r∈[r_ ,r_ ] and δ∈[δmin,δmax]δ∈[ _ , _ ], δmin>0 _ >0. The policy u∈=u:→u∈U=\u:X→A\ is assumed to be stationary and deterministic. Rewards rjj=1t\r_j\_j=1^t and durations δjj=1t\ _j\_j=1^t are drawn i.i.d. across steps from a fixed distribution given (x,a)(x,a). The objective is to maximize the relative value under the optimal reward rate: q∗(x,a)=[r(x,a)]−ρ∗[δ(x,a)],ρ∗=maxu∈ρ(u),q^*(x,a)\;=\;E [r(x,a) ]-ρ^*\,E [δ(x,a) ], ρ^*= _u∈Uρ(u), (5) so that a policy is optimal when it selects a∗=argmaxa∈q∗(x,a)a^*= _a∈Aq^*(x,a) at every context. When the reward rate is not optimal, ρ≠ρ∗ρ≠ρ^*, we denote the relative value as qρ(x,a)q_ρ(x,a). Regret is thus defined as Rt=ρ∗∑j=1tδj−∑j=1trj.R_t=ρ^* _j=1^t _j- _j=1^tr_j. (6) To minimize regret, the algorithm starts by estimating ρ∗ρ^* with ρ^t ρ_t using the past history (rj,δj)j=1t\(r_j, _j)\_j=1^t and performs an exhaustive search over all possible policies ρ^t=maxu∈ρ¯t(u)−bt(u),ρ¯t(u)=∑j=1taj−1=u(xj−1)rj∑j=1taj−1=u(xj−1)δj, ρ_t= _u∈U\, \ ρ_t(u)-b_t(u) \, ρ_t(u)= _j=1^t 1\a_j-1=u(x_j-1)\\,r_j _j=1^t 1\a_j-1=u(x_j-1)\\, _j, (7) which subtracts a confidence width bt(u)b_t(u), giving a lower bound satisfying ρ^t≤ρ∗ ρ_t≤ρ^* with high probability. Actions are then selected optimistically with the estimated relative value, where r¯t r_t and δ¯t δ_t are sample averages and ctc_t is a confidence bonus at=argmaxa∈r¯t(xt,a)−ρ^tδ¯t(xt,a)+ct(xt,a).a_t= _a∈A\; \ r_t(x_t,a)- ρ_t\, δ_t(x_t,a)+c_t(x_t,a) \. (8) This algorithm yields O(logT)O( T) expected regret bound but requires enumerating ||=|||||U|=|A|^|X| policies at each step. More details on the algorithm’s construction of ctc_t and bt(u)b_t(u) are in Section A.2. 3 Reward-rate Policy Gradient In this section, we propose RPG for optimizing the reward rate practically. We provide an outline for RPG in Section 3.1 that shows how we avoid previous algorithms’ limitations. Section 3.2 shares our theoretical insights, which were validated through experiments in Section 3.3. 3.1 Main Algorithm Current limitations. We first describe limitations of current approaches and then share the intuition behind our algorithm design. The SMDP policy gradient theorem in Equation 2 tells us it is viable to do gradient-based optimization on a parameterized policy to maximize reward rate. However, the core constraint in practice is that it requires multiple rollouts of the on-policy trajectory (rj,δj)j=1T\(r_j, _j)\_j=1^T to estimate ρ(π)ρ(π). Since the time δj _j for each step is variable and costly, this makes the estimation take a long time, at least ∑j≤Tδj _j≤ T _j, for each gradient update, making it impractical for training efficient MLE agents. On the other hand, the continuous-time bandit approach is an alternative by optimizing for the relative reward using an off-policy reward rate estimate ρ^t ρ_t (Equation 7), but the limitation being that it involves enumeration over all possible policies. Proposed algorithm. Given the limitations discussed above, our design principle for creating a practical reward-rate RL algorithm is as follows. We first use a predictor fρf_ρ fitted on the historical buffer ℬB to approximate the reward rate and use the predicted rate ρ ρ to calculate relative reward r~ r. This design allows us to avoid sampling on-policy long-horizon rollouts. Unlike the continuous-time bandit setting where policies are deterministic (György et al., 2007), LMs have stochastic policies, so we apply policy gradient updates on expected relative reward to train LMs agents. Our algorithm is shown in Algorithm 1. Ideally, the reward-rate predictor fρf_ρ should approximate the optimal ρ∗ρ^*, which then serves as ρ ρ for the relative reward. However, directly estimating ρ^≈ρ∗ ρ≈ρ^* is difficult, even in the bandit setting, since it requires enumeration over the policy space. Thus, to make the approach practical, we set fρf_ρ to estimate the greedy policy’s reward rate. This choice is supported by our results in the bandit setting (Sections 3.2 and 3.3), where we demonstrate that a softmax policy with NPG updates converges to the optimal ρ∗ρ^* through both theory and experiments. Empirically on MLE-Bench and NanoGPT, we show that PPO also performs well despite it being a first-order approximation of NPG. The estimator fρf_ρ can be anything that estimates the greedy reward rate, such as Bayesian inference, Robbins-Monro iteration, or neural networks. Algorithm 1 RPG 1: Policy πθ _θ, greedy reward rate estimator fρf_ρ, buffer ℬB, ρ^1←0 ρ_1← 0. 2: ⊳ fρf_ρ is a Normal-Inverse-Wishart (NIW) posterior, but it can also be Robbins-Monro or neural networks. ⊲ 3: for t=0t=0 to T−1T-1 do 4: Sample actions at∼πθt(⋅∣st)a_t _ _t(· s_t), execute ata_t and observe reward rt+1r_t+1 and time δt+1 _t+1. 5: Store ata_t, δt+1 _t+1 and rt+1r_t+1 in buffer ℬB. 6: Fit fρf_ρ on ℬB, then get ρ^t+1=fρ(ℬ) ρ_t+1=f_ρ(B). ⊳ Fit NIW posterior and draw samples. 7: Calculate relative reward r~t+1=rt+1−ρ^t+1δt+1 r_t+1=r_t+1- ρ_t+1\, _t+1 and advantage A^t+1 A_t+1. 8: θt+1←PPO(θt,A^t+1) _t+1← PPO( _t, A_t+1). 3.2 Theoretical Analysis Setup. We share our findings in the bandit setting, which provides motivation for approximating the greedy reward rate for ρ ρ and using PPO in our algorithm design. Our analysis is focused on the expected setting with a softmax policy πθ(a∣x)=exp(θa,x)/∑bexp(θb,x) _θ(a x)= ( _a,x)/ _b ( _b,x). Our assumptions are as follows: (1) boundedness in reward and time, (2) positive gap condition for relative value minx[q∗(x,a∗(x))−maxa≠a∗(x)q∗(x,a)]>0 _x [q^*(x,a^*(x))- _a≠ a^*(x)q^*(x,a) ]>0, and (3) contexts are drawn i.i.d. from p(x)p(x). Details of the analysis are in Appendix B. Approximate Dinkelbach iteration. We first analyze under which conditions an arbitrary sequence of reward rates ρkk=1∞\ _k\_k=1^∞ would be near the optimal ρ∗ρ^*. Theorem 1 shows that if ρk+1 _k+1 is within a asymmetric γ-tolerance band around the greedy reward rate ρ(uρk)ρ(u_ _k) at every step k≥2k≥ 2, then the sequence would be contained within an interval around the optimal reward rate ρk∈[(1−Rγ/(1+γ))ρ∗,(1+γ)ρ∗] _k∈[(1-Rγ/(1+γ))ρ^*,(1+γ)ρ^*] as k→∞k→∞, which suggests we are approximately doing a Dinkelbach iteration (Dinkelbach, 1967). When the estimation is exact, γ=0γ=0, the sequence reaches ρ∗ρ^* in finite steps. Theorem 1 suggests that we can use an estimation of the greedy reward rate as ρ ρ for relative reward calculation and iterate greedily to approximate ρ∗ρ^*. This finding holds tremendous value in practice, since predicting the greedy reward rate at the current ρk _k is an easier task than estimating ρ∗ρ^*, where the latter requires exhaustive search over the policy space but the former can be approximated with historical samples in ℬB by selecting top solutions according to r~(ρk)=r−ρkδ r( _k)=r- _k\,δ. Theorem 1. Let γ∈[0,1]γ∈[0,1] be a tolerance term, ρ1∈[0,ρmax] _1∈[0, _ ], and χ:=1−δmin/δmax∈(0,1)χ:=1- _ / _ ∈(0,1). For every k≥1k≥ 1, let uρk(x)∈argmaxaqρk(x,a)u_ _k(x)∈ _aq_ _k(x,a) be a greedy policy under any tie-breaking and ρ(uρk)ρ(u_ _k) is the reward rate of the greedy policy, let ρ~k+1 ρ_k+1 be any value with the sequence satisfying the condition ρ(uρk)/(1+γ)≤ρ~k+1≤(1+γ)ρ(uρk)ρ(u_ _k)/(1+γ)≤ ρ_k+1≤(1+γ)\,ρ(u_ _k), ρk+1=clip[0,ρmax](ρ~k+1) _k+1=clip_[0, _ ]( ρ_k+1) where clip[0,ρmax](z):=minmaxz,0,ρmaxclip_[0, _ ](z):= \ \z,0\,\, _ \, ρmax=rmax/δmin _ =r_ / _ . Then, for every k≥2k≥ 2, ρk∈[(1−Rγ/(1+γ))ρ∗−(χ/(1+γ))k−2ρmax,(1+γ)ρ∗], _k∈ [ (1-Rγ/(1+γ) )ρ^*- (χ/(1+γ) )^k-2 _ ,\,(1+γ)ρ^* ], (9) where R=δmax/δmin>1R= _ / _ >1. If γ=0γ=0, then the sequence increases monotonically ρ2≤ρ3≤⋯≤ρ∗ _2≤ _3≤·s≤ρ^* and the iteration reaches ρk=ρ∗ _k=ρ^* after finite steps. NPG and reward-rate RL. Theorem 1 suggests that if the predicted rate approximates the greedy rate, then it would be contained within an interval around the optimal ρ∗ρ^*, but it does not make any statements about the training dynamics, so we further analyze the behavior of πθ _θ. Define Aθ(ρ,x)=(qρ(x,⋅)−⟨πθ(⋅∣x),qρ(x,⋅)⟩ 1)A_θ(ρ,x)=(q_ρ(x,·)- _θ(· x),q_ρ(x,·) \, 1), where =[1…1]∈ℝ|| 1=[1… 1] ^|A|. The gradient update rules are θx′←θx+ηAθ(ρ,x)(NPG),θx′←θx+ηπθ(⋅∣x)⊙Aθ(ρ,x)(Standard PG).θ _x← _x+η\,A_θ(ρ,x) (NPG), θ _x← _x+η\, _θ(· x) A_θ(ρ,x) (Standard PG). We thus hypothesize that NPG is more suited for reward-rate RL since it has less coupling between πθ _θ and the magnitude of AθA_θ. Coupling would be undesirable since θx _x is updated proportionally to the policy weight πθ(⋅∣x) _θ(· x), reducing the signal for actions with large advantage AθA_θ but low πθ(⋅∣x) _θ(· x)11 1 In standard RL setting, this finding has already been observed Kakade (2001), but in reward-rate RL, AθA_θ is a function of ρ which is estimated from policy πθ _θ, so the amount of coupling is stronger.. We establish the validity of NPG updates in Theorem 2. It shows that the natural gradient updates are equivalent to a Boltzmann policy on the relative value qρ¯k(x,a)q_ ρ_k(x,a) with a running average rate ρ¯k ρ_k. If ρ1 _1 is initialized within [0,ρ∗][0,ρ^*], and greedy reward rate estimation is exact ρk+1=ρ(uρ¯k) _k+1=ρ(u_ ρ_k), then the error between ρ¯k ρ_k and ρ∗ρ^* shrinks in O(k−(1−χ))O(k^-(1-χ)). Furthermore, when ρ¯k ρ_k is within ε¯ of ρ∗ρ^*, which Equation 10 guarantees after finitely many steps, the greedy action at ρ¯k ρ_k becomes the optimal one, and the policy increases the optimal arm’s probability towards 11. This finding suggests that NPG methods bring performance benefits in reward-rate RL and motivates us to apply PPO in practice, a first-order approximation of NPG (Schulman et al., 2017). Theorem 2. Let ρ1,ρ2,… _1, _2,… be a sequence of reward rates and let θ(k)θ^(k) be the softmax logits after k NPG update steps with learning rate η from the zero initialization, with ρ¯k=∑n≤kρn/k ρ_k= _n≤ k _n/k. Then for every x and k≥1k≥ 1, πθ(k) _θ^(k) is the Boltzmann policy of qρ¯k(x,a)q_ ρ_k(x,a) with inverse temperature ηkη k. If moreover ρ1∈[0,ρ∗] _1∈[0,ρ^*] and ρk+1=ρ(uρ¯k) _k+1=ρ(u_ ρ_k), then ρ¯k ρ_k increases to ρ∗ρ^* in O(k−(1−χ))O(k^-(1-χ)) ρ∗−ρ¯k≤(ρ∗−ρ¯1)exp(1−χ)(k+1)−(1−χ).ρ^*- ρ_k\;≤\;(ρ^*- ρ_1)\, (1-χ)\,(k+1)^-(1-χ). (10) Furthermore, let δspr∈(0,δmax−δmin] _spr∈(0, _ - _ ] be the maximum time gap, and Δ∗ ^* be the minimum relative value gap under ρ∗ρ^*, define ε¯:=Δ∗/(2δspr) := ^*/(2 _spr), for every k≥k0k≥ k_0, the greedy policy is optimal πθ(k)(a∗∣x)≥1−(||−1)exp(−ηΔ∗k/2),k0:=⌈(exp(1−χ)(ρ∗−ρ¯1)/ε¯)1/(1−χ)⌉. _θ^(k)(a^* x)≥ 1-(|A|-1)\, (-η\, ^*k/2 ), k_0:= ( (1-χ)\,(ρ^*- ρ_1)/ )^1/(1-χ) . (11) 3.3 Bandit Validation (a) 4 contexts ||=4|X|=4. (b) 64 contexts ||=64|X|=64. Figure 2: Our proposed approach NPG-NIW obtains similar or lower regret than C-UCB while avoiding exhaustive search over policies. Standard policy gradient (SPG-NIW) also performs worse than NPG, verifying our theory that NPG’s decoupling helps reward-rate RL. y-axis is regret at the final step over 100100 trials, x-axis is the number of arms K, and 4 environments are considered: (E1) independent reward and time, (E2) time and reward have 0.80.8 correlation, (E3) time and reward have −0.8-0.8 correlation, and (E4) reward and time are lognormal with +0.5+0.5 correlation. Our theory suggests that NPG has a decoupling effect on reward-rate optimization, and using greedy reward rate as ρ ρ can also approximate the optimal rate as iterations grow. The former brings performance benefits, whereas the latter is a feasibility benefit in large action and context spaces where an exhaustive search over the policy space is intractable. We conduct experiments to examine whether these benefits hold true in practice. Our finding is in Figure 2, and NPG-NIW uses softmax policy with NPG updates and NIW posterior fitted on greedy historical samples for ρ ρ estimation. Continuous-time upper confidence bound (C-UCB) is the approach by György et al. (2007). We first observe that if we use C-UCB’s constants exactly as in the paper, it breaks since the constants are too large (C-UCB, theory), but treating them as tunable hyperparameters allows it to work (C-UCB, tuned). When only having small context and action spaces (Figure 2(a)), C-UCB uses exhaustive search for ρ ρ estimation. When the context and action space is large (Figure 2(b)), C-UCB cannot afford to enumerate over all policies, so we use an approximation where 2048 policies are drawn uniformly at random from past visited ones to predict ρ ρ. Across both settings, we observe that NPG-NIW obtains lower or similar regret as C-UCB baselines, thus confirming our hypothesis that greedy reward rate estimation can be used for ρ ρ without enumerating the policy space. When comparing NPG with standard policy gradient updates (SPG-NIW), we observe that it consistently obtains higher regret than NPG, supporting our theoretical finding that NPG provides performance benefits. More experiments and implementation details are in Sections C.1, D.1 and D.2. 4 Reward-rate for MLE Agent In this section, we present our results demonstrating that RPG finds good solutions in less time than vanilla RL, suggesting that reward-rate RL is a practical solution for training efficient agents. We provide our empirical results on MLE-Bench and NanoGPT in Sections 4.2 and 4.3. 4.1 Instantiating Language Models for Reward-rate Policies Environment setup. The theory in Section 3.2 suggests that we should approximate the greedy reward rate for ρ ρ and use PPO for training, but training MLE agents also requires deliberate environment setup, which we share in this section. We start with a base prompt as our initial state s0s_0 containing task descriptions such as data loading path, meaning of features, and the performance metric used for the evaluation. We also explicitly ask the LM agent to take efficiency into consideration by encouraging it to use a subset of features. An action a0∼π(⋅∣s0)a_0 π(· s_0) is the generated script by the agent, and the code is executed in a sandbox environment to produce time δ1 _1 and reward r1r_1. If the script successfully executes, r1r_1 would be a performance score. A timeout budget δmax _ is set as the upper limit of the time allowed for script execution. For scripts that fail to run, such as when they have a bug, we encounter trouble assigning their reward rate since their reward is 00 and r~=−ρ^δ r=- ρ\,δ is maximized by δ→0δ→ 0. This implies a script that fails on its first line dominates all honest attempts. Thus, we charge every failure the full time limit δmax _ and a constant penalty of −10-10, preventing the policy from reward-hacking by failing fast, and executable scripts have higher rewards than non-executable ones. State transition with self-improvement. Besides setting up the environment and reward properly, we also want our agent to improve based on its own previous attempts, so we would ideally have a self-improving agent making progress on the task over time. We discuss how we connect the agent’s own solutions to form self-improvement loops. At every step, we store the observed transition (at,rt+1,δt+1)(a_t,r_t+1, _t+1) into the buffer ℬB. The next state st+1s_t+1 is constructed by concatenating the base prompt s0s_0 with the previous code a′a and a self-improvement phrase i0i_0 explicitly asking the agent to improve upon the performance score r′r and time δ′δ . The selection of (a′,r′,δ′)(a ,r ,δ ) from buffer ℬB is based on a top-k selection according to the relative reward r~(ρ^t+1)=r′−ρ^t+1δ′ r( ρ_t+1)=r - ρ_t+1\,δ . This naturally gives a self-improvement chain where the agent is tasked with making progress on its best proposed solutions observed thus far. Reward-rate predictor. At each step t, our agent uses its history of attempts to predict the reward rate ρ ρ. The reward-rate estimator fρf_ρ is a NIW posterior with joint Gaussian likelihood (Murphy, 2012) fitted from historical samples observed in buffer ℬB. We model the pair as (racc,logδacc)(r_acc, _acc), where raccr_acc and δacc _acc are the accumulated reward and time, and we take the logarithm because we need time to be positive. We fit the posterior at each step, and sample M reward rates from it and take the 95th percentile ρ^t+1=Quantile0.95(ρmm=1M) ρ_t+1=Quantile_0.95(\ _m\_m=1^M). More details are in Section C.2. Table 1: RPG vs. vanilla RL on 22 MLE-Bench Lite tasks, evaluated on mean performance score across self-improving chains within a fixed-time budget. (↑)( ) and (↓)( ) means higher and lower is better, respectively. %\% improv. is the relative improvement of RPG over vanilla. On average, RPG performs 19.2% better than vanilla, after dropping the highest and lowest numbers. Task RPG Vanilla % improv. Task RPG Vanilla % improv. ranzcr-clip (↑)( ) 0.089 ± 0.046 0.014 ± 0.014 ++530.5 plant-pathology (↑)( ) 0.526 ± 0.048 0.498 ± 0.047 ++5.6 aptos2019 (↑)( ) 0.564 ± 0.046 0.315 ± 0.120 ++78.9 pizza (↑)( ) 0.760 ± 0.008 0.728 ± 0.019 ++4.4 mlsp-birds (↑)( ) 0.333 ± 0.154 0.226 ± 0.120 ++46.9 tabular-2022 (↑)( ) 0.840 ± 0.013 0.806 ± 0.028 ++4.2 siim-isic (↑)( ) 0.660 ± 0.021 0.456 ± 0.112 ++44.8 whale (↑)( ) 0.493 ± 0.005 0.749 ± 0.091 −-34.2 textnorm-en (↑)( ) 0.880 ± 0.048 0.714 ± 0.068 ++23.2 leaf-classif (↓)( ) 0.259 ± 0.063 0.564 ± 0.059 ++54.0 textnorm-ru (↑)( ) 0.889 ± 0.021 0.730 ± 0.022 ++21.7 spooky-author (↓)( ) 0.440 ± 0.015 0.690 ± 0.003 ++36.2 tabular-2021 (↑)( ) 0.907 ± 0.010 0.778 ± 0.046 ++16.5 nyc-taxi (↓)( ) 4.309 ± 0.029 5.349 ± 0.387 ++19.4 insults (↑)( ) 0.834 ± 0.025 0.761 ± 0.024 ++9.6 dog-breed (↓)( ) 4.786 ± 0.001 4.924 ± 0.139 ++2.8 jigsaw-toxic (↑)( ) 0.948 ± 0.009 0.868 ± 0.028 ++9.2 dogs-vs-cats (↓)( ) 0.688 ± 0.005 0.703 ± 0.001 ++2.1 histopathologic (↑)( ) 0.594 ± 0.107 0.554 ± 0.021 ++7.2 nomad2018 (↓)( ) 0.064 ± 0.002 0.064 ± 0.002 ++0.3 aerial-cactus (↑)( ) 0.954 ± 0.018 0.893 ± 0.010 ++6.8 denoising-docs (↓)( ) 0.094 ± 0.002 0.085 ± 0.013 −-10.1 Figure 3: Example of generated scripts by RPG and vanilla RL. RPG proposes scripts that use time more efficiently by training on one type of feature with logistic regression, whereas vanilla RL chooses to train on two types of features with 1200 gradient-boosted trees. 4.2 Evaluation Results on MLE-Bench We provide our results for MLE-Bench in Table 1. We use Qwen3.5-4B as our LM and apply the method in Section 4.1. Our baseline is vanilla RL, where rewards are performance scores only and time budget is δmax _ . We run RPG and the baseline until convergence. Our evaluation metric is the mean reward obtained by the self-improvement chains at every time budget. Specifically, at every step, there are batch_size self-improvement chains, and for each chain, a time δ and reward r are earned. We calculate the mean reward r across chains at every cumulative time budget t∈[0,200,…,T]t∈[0,200,…,T] with 200200-second intervals. The final time T is decided by taking the cumulative time containing 95% of the self-improvement chains. Each task is evaluated with three repetitions. Across the 2222 tasks RPG show an average 19.2% improvement over vanilla at the final time budget, after dropping the task with highest and lowest improvements. The two methods’ proposed MLE scripts have distinct differences in wall clock time. An example of the budget curve along with the generated script is shown in Figure 3 for the spooky-author task. Vanilla uses a concatenation of different features and trains a gradient-boosted tree in 13 seconds with log loss 0.727, whereas RPG uses count features along with logistic regression to obtain 0.469 log loss in 5 seconds. The budget curve behind every task is provided in Section D.3, and implementation details are in Section C.3. 4.3 Evaluation Results on NanoGPT Figure 4: Budget curve and an example self-improvement step on NanoGPT. (Left) At the final budget, RPG improves over vanilla by 85.7%85.7\%. (Right) Example of a self-improvement edit by RPG, which rewrites Newton–Schulz in the Muon optimizer as a triton kernel with a compiled iteration and adds a learning-rate warmup. The new edit crosses the loss target in less time. Motivation and setup. The MLE-Bench setting discussed above is based on Kaggle competitions. To extend the scope of reward-rate RL agents, we also apply RPG to the NanoGPT speedrun challenge (Karpathy, 2022; Jordan et al., 2024), where the task is to train a tiny LM down to a validation cross-entropy loss of 3.283.28 on FineWeb (Penedo et al., 2024) with eight GPUs as fast as possible. This task makes the time taken to reach the 3.283.28 target loss a natural choice to optimize for. Thus, as a baseline, we implement vanilla RL using time to reach the 3.283.28 target loss as reward, and compare with RPG using relative reward as the learning signal. Setup. Both vanilla and RPG have a linear bonus terms for crossing the 3.283.28 target loss, w(3.28−L)w\,(3.28-L). NanoGPT requires writing tens of thousands of tokens in its code. Thus, to ensure efficient training and focus on making modifications rather than writing executable code, we ask the agent to read prompt sts_t and propose high-level delta code changes along with their reasoning for the modification, which is the action ata_t. A separate, independent, frontier model with deterministic output (zero temperature) serves as the implementer, which reads the action ata_t proposed by the agent and implements the changes. More details are in Section C.4. Results. We present our results on NanoGPT in Figure 4, where we observe that RPG gives lower validation loss given the same time budget than vanilla RL. At the final time budget, we observe that RPG has an 85.7% margin below 3.283.28 compared to the baseline. This can be explained by its objective of using the reward rate, which intrinsically optimizes for the loss per unit of time. An example of a self-improvement edit by RPG is shown on the right side of Figure 4, where the agent writes a triton kernel, a compiled iteration, and a learning rate warm-up in its parent script to drive the loss to the target in less time. 5 Related work Machine learning engineering agents. Many works currently ask whether an agent can carry out machine learning engineering tasks in an end-to-end fashion (Nam et al., 2026; Tang et al., 2023; Chan et al., 2025; Li et al., 2024; Zhang et al., 2026; Qiang et al., 2026). Similar lines of work also include using agents on data science tasks such as data preprocessing and hyperparameter tuning (Grosnit et al., 2024; Bendinelli et al., 2025; Zhang et al., 2023; Liu et al., 2024). Most agentic approaches in these tasks involve scaffolding a frozen model, such as by searching the solution space as a tree (Jiang et al., 2025), which do not train MLE-specific experience into the model weights. One of the closest works to ours is by Yang et al. (2025), who trained a small LM with RL on MLE-Bench and observed that actions take a variable amount of time, but the focus was on mitigating the time bias favoring faster solutions in asynchronous RL updates. By contrast, our work directly treats time as a cost term in the RL objective using the SMDP reward-rate formulation and trains agents to learn to take efficient actions. Reward-rate RL. We acknowledge that the reward-rate objective itself is classical. The SMDP formulation provides the theoretical framework for our method (Sutton et al., 1999b), and the continuous-time bandit (György et al., 2007) gives the idea of using off-policy estimates for the optimal reward rate. Similar problems also appear in renewal processes with bandit feedback (Cayci et al., 2019) where the objective is to maximize reward within a fixed horizon. A core mathematical toolkit is Dinkelbach iteration for fractional programming (Dinkelbach, 1967; Schaible, 1976), where all of these works aforementioned rely on it to transform a ratio objective into a sequence of linear ones. Budgeted and knapsack bandits (Badanidiyuru et al., 2013; Ding et al., 2013; Xia et al., 2016) also consider actions with costs, but they enforce a hard constraint on a finite budget. In RL for LMs literature, the environment is episodic with actions assumed to have a constant cost (Schulman et al., 2017; DeepSeek-AI et al., 2025). A survey on agentic RL with LMs (Zhang et al., 2025a) also uses the same formulation, with no notion of how long an action can take. Guan et al. (2025) trades off latency against cost, but requires exploring the Pareto frontier in a multi-agent framework. To our knowledge, no existing work applies reward-rate RL in modern LM agents. Self-improvement. Performing self-improvement by asking an agent to revise its own previous attempts is a well-established pattern, such as critiquing its outputs directly (Madaan et al., 2023), reflecting on feedback from the environment (Shinn et al., 2023), training on the attempts that happened to succeed (Zelikman et al., 2022), and rewriting the scaffolding (Zelikman et al., 2023). The idea of self-improvement loops has also been explored in discovery tasks (Yuksekgonul et al., 2026). Closer to our setting, the loop can be driven by execution feedback, where a model repairs its own program from the error message that the program produced (Chen et al., 2023). The loop can also be used on an agent archive to create a tree of diverse and capable agents through evolutionary search (Zhang et al., 2025b). Zeng et al. (2025) proposes using a balance score for the exploration-exploitation trade-off and updating models on samples that are above a score threshold. Our work focuses on using self-improving loops as state transitions to train agents. 6 Conclusion We proposed a practical reward-rate RL algorithm, namely RPG, for training efficient MLE agents using the SMDP formulation, where each action takes variable time, and the objective is long-run reward per unit of time. RPG uses greedy policy’s reward rate estimates for the relative reward calculation and directly maximizes it through PPO. It avoids enumeration over all policies and time-costly on-policy rollouts, two limitations for directly applying the SMDP methods in practice. Through theoretical analysis, simulation experiments, and empirical experiments on MLE-Bench and NanoGPT, we demonstrated that RPG produces an agent policy that learns to take actions efficiently and maximizes reward. There are some limitations in this work. The scope of the empirical experiments is focused on MLE tasks, but other agentic RL tasks where actions cost variable time exist, so it would be interesting to see how reward-rate RL does on these tasks. Theoretically, our analysis focuses on the population-level setting in expectation. A finite-sample analysis could provide further insight. AI Use Statement In this work, we used generative AI tools for: implement methods, support qualitative and thematic data analysis, process experiment data, assist in the writing of proofs, and design or provide feedback on research methodology or experiments. We have not used generative AI tools for: generate synthetic data sets, help develop theoretical models or conceptual frameworks, provide critical ingredients for proving mathematical claims, formulate mathematical claims, propose or refine hypotheses, assist with translation, clean and reformat dataset, interpret results. Additionally, we used generative AI tools for: create or modify scientific figures or images, create or edit software code, creation of artifacts, draft parts of a research paper, summarize or analyse existing literature. We have reviewed all AI-assisted work. We have checked through all the code edited by AI, and all experiments involving AI were verified. We take responsibility for the final content of this work, including text, claims or artifacts produced with the aid of generative AI. References Amari (1998) S. Amari Natural Gradient Works Efficiently in Learning. Neural Computation 10 (2), p. 251–276. Cited by: §1. Auer et al. (2002) P. Auer, N. Cesa-Bianchi, and P. Fischer Finite-time Analysis of the Multiarmed Bandit Problem. Machine Learning 47 (2), p. 235–256 (en). External Links: Document Cited by: Assumption 2. Badanidiyuru et al. (2013) A. Badanidiyuru, R. Kleinberg, and A. Slivkins Bandits with Knapsacks. In 2013 IEEE 54th Annual Symposium on Foundations of Computer Science, p. 207–216. Note: arXiv:1305.2545 [cs]Comment: An extended abstract of this work has appeared in the 54th IEEE Symposium on Foundations of Computer Science (FOCS 2013). 55 pages. Compared to the initial "full version" from May’13, this version has a significantly revised presentation and reflects the current status of the follow-up work. Also, this version contains a stronger regret bound in one of the main results External Links: Link, Document Cited by: §5. Bendinelli et al. (2025) T. Bendinelli, A. Dox, and C. Holz Exploring llm agents for cleaning tabular machine learning datasets. arXiv preprint arXiv:2503.06664. Cited by: §5. Cayci et al. (2019) S. Cayci, A. Eryilmaz, and R. Srikant Learning to Control Renewal Processes with Bandit Feedback. Proceedings of the ACM on Measurement and Analysis of Computing Systems 3 (2), p. 1–32 (en). External Links: ISSN 2476-1249, Link, Document Cited by: §5. Chan et al. (2025) J. S. Chan, N. Chowdhury, O. Jaffe, J. Aung, D. Sherburn, E. Mays, G. Starace, K. Liu, L. Maksin, T. Patwardhan, et al. Mle-bench: evaluating machine learning agents on machine learning engineering. In International Conference on Learning Representations, Vol. 2025, p. 50466–50494. Cited by: §1, §5. Chen et al. (2023) X. Chen, M. Lin, N. Schärli, and D. Zhou Teaching Large Language Models to Self-Debug. arXiv. Note: arXiv:2304.05128 Cited by: §5. Christiano et al. (2017) P. Christiano, J. Leike, T. B. Brown, M. Martic, S. Legg, and D. Amodei Deep Reinforcement Learning from Human Preferences. In Advances in Neural Information Processing Systems, Note: arXiv:1706.03741 Cited by: §1. Das et al. (1999) T. K. Das, A. Gosavi, S. Mahadevan, and N. Marchalleck Solving Semi-Markov Decision Problems Using Average Reward Reinforcement Learning. Management Science 45 (4), p. 560–574. External Links: ISSN 0025-1909, Link Cited by: §A.1, §1, §2.1. DeepSeek-AI et al. (2025) DeepSeek-AI, D. Guo, D. Yang, H. Zhang, J. Song, P. Wang, Q. Zhu, R. Xu, R. Zhang, S. Ma, X. Bi, X. Zhang, X. Yu, Y. Wu, Z. F. Wu, Z. Gou, Z. Shao, Z. Li, Z. Gao, A. Liu, B. Xue, B. Wang, B. Wu, B. Feng, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, D. Dai, D. Chen, D. Ji, E. Li, F. Lin, F. Dai, F. Luo, G. Hao, G. Chen, G. Li, H. Zhang, H. Bao, H. Xu, H. Wang, H. Ding, H. Xin, H. Gao, H. Qu, H. Li, J. Guo, J. Li, J. Wang, J. Chen, J. Yuan, J. Qiu, J. Li, J. L. Cai, J. Ni, J. Liang, J. Chen, K. Dong, K. Hu, K. Gao, K. Guan, K. Huang, K. Yu, L. Wang, L. Zhang, L. Zhao, L. Wang, L. Zhang, L. Xu, L. Xia, M. Zhang, M. Zhang, M. Tang, M. Li, M. Wang, M. Li, N. Tian, P. Huang, P. Zhang, Q. Wang, Q. Chen, Q. Du, R. Ge, R. Zhang, R. Pan, R. Wang, R. J. Chen, R. L. Jin, R. Chen, S. Lu, S. Zhou, S. Chen, S. Ye, S. Wang, S. Yu, S. Zhou, S. Pan, S. S. Li, S. Zhou, S. Wu, S. Ye, T. Yun, T. Pei, T. Sun, T. Wang, W. Zeng, W. Zhao, W. Liu, W. Liang, W. Gao, W. Yu, W. Zhang, W. L. Xiao, W. An, X. Liu, X. Wang, X. Chen, X. Nie, X. Cheng, X. Liu, X. Xie, X. Liu, X. Yang, X. Li, X. Su, X. Lin, X. Q. Li, X. Jin, X. Shen, X. Chen, X. Sun, X. Wang, X. Song, X. Zhou, X. Wang, X. Shan, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. Zhang, Y. Xu, Y. Li, Y. Zhao, Y. Sun, Y. Wang, Y. Yu, Y. Zhang, Y. Shi, Y. Xiong, Y. He, Y. Piao, Y. Wang, Y. Tan, Y. Ma, Y. Liu, Y. Guo, Y. Ou, Y. Wang, Y. Gong, Y. Zou, Y. He, Y. Xiong, Y. Luo, Y. You, Y. Liu, Y. Zhou, Y. X. Zhu, Y. Xu, Y. Huang, Y. Li, Y. Zheng, Y. Zhu, Y. Ma, Y. Tang, Y. Zha, Y. Yan, Z. Z. Ren, Z. Ren, Z. Sha, Z. Fu, Z. Xu, Z. Xie, Z. Zhang, Z. Hao, Z. Ma, Z. Yan, Z. Wu, Z. Gu, Z. Zhu, Z. Liu, Z. Li, Z. Xie, Z. Song, Z. Pan, Z. Huang, Z. Xu, Z. Zhang, and Z. Zhang DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Nature 645 (8081), p. 633–638. Note: arXiv:2501.12948 [cs] External Links: ISSN 0028-0836, 1476-4687, Link, Document Cited by: §5. Ding et al. (2013) W. Ding, T. Qin, X. Zhang, and T. Liu Multi-Armed Bandit with Budget Constraint and Variable Costs. Proceedings of the AAAI Conference on Artificial Intelligence 27 (1), p. 232–238 (en). External Links: ISSN 2374-3468, Link, Document Cited by: §5. Dinkelbach (1967) W. Dinkelbach On Nonlinear Fractional Programming. Management Science 13 (7), p. 492–498. External Links: ISSN 0025-1909, Link Cited by: §B.2, §1, §3.2, §5. Grosnit et al. (2024) A. Grosnit, A. Maraval, J. Doran, G. Paolo, A. Thomas, R. S. H. N. Beevi, J. Gonzalez, K. Khandelwal, I. Iacobacci, A. Benechehab, et al. Large language models orchestrating structured reasoning achieve kaggle grandmaster level. arXiv preprint arXiv:2411.03562 3 (4). Cited by: §5. Guan et al. (2025) Y. Guan, Q. Lan, S. Fei, D. Ding, D. Acharya, C. Wang, W. Y. Wang, and W. Hua Dynamic Speculative Agent Planning. arXiv. Note: arXiv:2509.01920 [cs.AI] External Links: Link, Document Cited by: §5. György et al. (2007) A. György, L. Kocsis, I. Szabó, and C. Szepesvári Continuous Time Associative Bandit Problems. In Proceedings of the 20th International Joint Conference on Artificial Intelligence, p. 830–835 (en). Cited by: §A.2, §A.2, §C.1.1, §1, §1, §2.2, §3.1, §3.3, §5, Assumption 3. Hayes et al. (2022) C. F. Hayes, R. Rădulescu, E. Bargiacchi, J. Källström, M. Macfarlane, M. Reymond, T. Verstraeten, L. M. Zintgraf, R. Dazeley, F. Heintz, E. Howley, A. A. Irissappane, P. Mannion, A. Nowé, G. Ramos, M. Restelli, P. Vamplew, and D. M. Roijers A practical guide to multi-objective reinforcement learning and planning. Autonomous Agents and Multi-Agent Systems 36 (1), p. 26. External Links: ISSN 1387-2532, 1573-7454, Link, Document Cited by: §1. Jagannathan (1966) R. Jagannathan On Some Properties of Programming Problems in Parametric form Pertaining to Fractional Programming. Management Science 12 (7), p. 609–615. External Links: Document Cited by: §B.2. Jiang et al. (2025) Z. Jiang, D. Schmidt, D. Srikanth, D. Xu, I. Kaplan, D. Jacenko, and Y. Wu AIDE: AI-Driven Exploration in the Space of Code. arXiv. Note: arXiv:2502.13138 Cited by: §1, §5. Jordan et al. (2024) K. Jordan, J. Bernstein, B. Rappazzo, et al. Modded-NanoGPT: Speedrunning the NanoGPT baseline. Note: https://github.com/KellerJordan/modded-nanogptGitHub repository Cited by: §C.4.1, §4.3. Kakade (2001) S. M. Kakade A Natural Policy Gradient. In Advances in Neural Information Processing Systems, Vol. 14. Cited by: §1, footnote 1. Karpathy (2022) A. Karpathy nanoGPT: The simplest, fastest repository for training/finetuning medium-sized GPTs. Note: https://github.com/karpathy/nanoGPTGitHub repository Cited by: §4.3. Kingma and Ba (2015) D. P. Kingma and J. Ba Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations, Note: arXiv:1412.6980 Cited by: §2.1. Li et al. (2024) Z. Li, Q. Zang, D. Ma, J. Guo, T. Zheng, M. Liu, X. Niu, Y. Wang, J. Yang, J. Liu, et al. Autokaggle: a multi-agent framework for autonomous data science competitions. arXiv preprint arXiv:2410.20424. Cited by: §5. Liu et al. (2024) S. Liu, C. Gao, and Y. Li Large language model agent for hyper-parameter optimization. arXiv preprint arXiv:2402.01881. Cited by: §5. Madaan et al. (2023) A. Madaan, N. Tandon, P. Gupta, S. Hallinan, L. Gao, S. Wiegreffe, U. Alon, N. Dziri, S. Prabhumoye, Y. Yang, S. Gupta, B. P. Majumder, K. Hermann, S. Welleck, A. Yazdanbakhsh, and P. Clark Self-Refine: Iterative Refinement with Self-Feedback. arXiv. Note: arXiv:2303.17651 Cited by: §5. Murphy (2012) K. P. Murphy Machine Learning: A Probabilistic Perspective. MIT Press. Cited by: §C.2, §4.1. Nam et al. (2026) J. Nam, J. Yoon, J. Chen, J. Shin, S. Arik, and T. Pfister Mle-star: machine learning engineering agent via search and targeted refinement. Advances in Neural Information Processing Systems 38, p. 116692–116712. Cited by: §1, §5. Penedo et al. (2024) G. Penedo, H. Kydlíček, L. Ben allal, A. Lozhkov, M. Mitchell, C. Raffel, L. Von Werra, and T. Wolf The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale. arXiv. Note: arXiv:2406.17557 Cited by: §4.3. Qiang et al. (2026) R. Qiang, Y. Zhuang, Y. Li, D. Sagar VK, R. Zhang, C. Li, I. Wong, S. Yang, P. Liang, C. Zhang, et al. Mle-dojo: interactive environments for empowering llm agents in machine learning engineering. Advances in Neural Information Processing Systems 38. Cited by: §5. Roijers et al. (2013) D. M. Roijers, P. Vamplew, S. Whiteson, and R. Dazeley A Survey of Multi-Objective Sequential Decision-Making. Journal of Artificial Intelligence Research 48, p. 67–113. Note: arXiv:1402.0590 Cited by: §1. Schaible (1976) S. Schaible Fractional Programming. I, on Dinkelbach’s Algorithm. Management Science 22 (8), p. 868–873. External Links: ISSN 0025-1909, Link Cited by: §B.2, §1, §5. Schulman et al. (2015) J. Schulman, S. Levine, P. Moritz, M. Jordan, and P. Abbeel Trust Region Policy Optimization. In Proceedings of the 32nd International Conference on Machine Learning, F. Bach and D. Blei (Eds.), Proceedings of Machine Learning Research, Vol. 37, p. 1889–1897 (en). Cited by: §1. Schulman et al. (2017) J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov Proximal Policy Optimization Algorithms. arXiv. Note: arXiv:1707.06347 [cs] External Links: Link, Document Cited by: §3.2, §5. Shinn et al. (2023) N. Shinn, F. Cassano, E. Berman, A. Gopinath, K. Narasimhan, and S. Yao Reflexion: Language Agents with Verbal Reinforcement Learning. arXiv. Note: arXiv:2303.11366 Cited by: §5. Sutton et al. (1999a) R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour Policy Gradient Methods for Reinforcement Learning with Function Approximation. In Advances in Neural Information Processing Systems, Vol. 12. External Links: Link Cited by: §1. Sutton et al. (1999b) R. S. Sutton, D. Precup, and S. Singh Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning. Artificial Intelligence 112 (1-2), p. 181–211 (en). External Links: ISSN 00043702, Link, Document Cited by: §A.1, §1, §2.1, §5. Tang et al. (2023) X. Tang, Y. Liu, Z. Cai, Y. Shao, J. Lu, Y. Zhang, Z. Deng, H. Hu, K. An, R. Huang, et al. Ml-bench: evaluating large language models and agents for machine learning tasks on repository-level code. arXiv preprint arXiv:2311.09835. Cited by: §5. Xia et al. (2016) Y. Xia, W. Ding, X. Zhang, N. Yu, and T. Qin Budgeted Bandit Problems with Continuous Random Costs. In Asian Conference on Machine Learning, G. Holmes and T. Liu (Eds.), Proceedings of Machine Learning Research, Vol. 45, Hong Kong, p. 317–332 (en). Cited by: §5. Yang et al. (2025) S. Yang, J. He-Yueya, and P. Liang Reinforcement Learning for Machine Learning Engineering Agents. arXiv. Note: arXiv:2509.01684 [cs] External Links: Link, Document Cited by: §1, §5. Yuksekgonul et al. (2026) M. Yuksekgonul, D. Koceja, X. Li, F. Bianchi, J. McCaleb, X. Wang, J. Kautz, Y. Choi, J. Zou, C. Guestrin, and Y. Sun Learning to Discover at Test Time. arXiv (en). Note: Version Number: 2Other Code: https://github.com/test-time-training/discover External Links: Link, Document Cited by: §5. Zelikman et al. (2023) E. Zelikman, E. Lorch, L. Mackey, and A. T. Kalai Self-Taught Optimizer (STOP): Recursively Self-Improving Code Generation. arXiv. Note: arXiv:2310.02304 Cited by: §5. Zelikman et al. (2022) E. Zelikman, Y. Wu, J. Mu, and N. D. Goodman STaR: Bootstrapping Reasoning With Reasoning. arXiv. Note: arXiv:2203.14465 Cited by: §5. Zeng et al. (2025) W. Zeng, Y. Huang, L. Zhao, Y. Wang, Z. Shan, and J. He B-STaR: monitoring and balancing exploration and exploitation in self-taught reasoners. In International Conference on Learning Representations, Cited by: §5. Zhang et al. (2025a) G. Zhang, H. Geng, X. Yu, Z. Yin, Z. Zhang, Z. Tan, H. Zhou, Z. Li, X. Xue, Y. Li, et al. The landscape of agentic reinforcement learning for llms: a survey. arXiv preprint arXiv:2509.02547. Cited by: §1, §5. Zhang et al. (2025b) J. Zhang, S. Hu, C. Lu, R. Lange, and J. Clune Darwin Gödel Machine: Open-Ended Evolution of Self-Improving Agents. arXiv. Note: arXiv:2505.22954 Cited by: §5. Zhang et al. (2023) S. Zhang, C. Gong, L. Wu, X. Liu, and M. Zhou Automl-gpt: automatic machine learning with gpt. arXiv preprint arXiv:2305.02499. Cited by: §5. Zhang et al. (2026) Y. Zhang, M. Khalifa, S. Bhushan, G. Murphy, L. Logeswaran, J. Kim, M. Lee, H. Lee, and L. Wang MLRC-bench: can language agents solve machine learning research challenges?. Advances in Neural Information Processing Systems 38. Cited by: §5. Appendix Appendix A Derivation A.1 Semi-MDP Policy Gradient Theorem Setup and notation. We consider a finite SMDP (Sutton et al., 1999b; Das et al., 1999) ℳ=(,,p,δmin,δmax)M=(S,A,p, _ , _ ) with finite state space S and finite action space A. At each decision epoch the agent observes a state s∈s∈S, selects an action a∈a∈A, and the environment returns a triple (r,δ,s′)(r,δ,s ) drawn from the transition kernel p(s′,r,δ∣s,a),p(s′∣s,a)=∫p(s′,r,δ∣s,a)drdδ,p(s ,r,δ s,a), p(s s,a)\;=\; \!\! p(s ,r,δ s,a)\,drdδ, (12) where r∈[rmin,rmax]r∈[r_ ,r_ ] is the reward accumulated over the transition and δ∈[δmin,δmax]δ∈[ _ , _ ] is its time (or duration). We write r¯(s,a):=[r∣s,a] r(s,a):=E[r s,a] and δ¯(s,a):=[δ∣s,a] δ(s,a):=E[δ s,a] for the conditional means under Equation 12. Crucially, the kernel does not depend on the policy parameters, so ∇θr¯(s,a)=0 _θ r(s,a)=0 and ∇θδ¯(s,a)=0 _θ δ(s,a)=0. Let πθ(a∣s) _θ(a s) be a differentiable stochastic policy parameterized by θ∈ℝdθ ^d. We drop the subscript θ for brevity when clear from context. We abbreviate, for a fixed policy π, δ¯π(s):=π[δ∣s]=∑aπ(a∣s)δ¯(s,a). δ^π(s)\;:=\;E_π[δ s]\;=\; _aπ(a s)\, δ(s,a). (13) Let pπ(s′∣s):=∑aπ(a∣s)p(s′∣s,a)p^π(s s):= _aπ(a s)\,p(s s,a) denote the transition matrix of the Markov chain. We write dπd^π for its unique stationary distribution, dπ(s′)=∑sdπ(s)∑aπ(a∣s)p(s′∣s,a)∀s′∈.d^π(s )\;=\; _sd^π(s) _aπ(a s)\,p(s s,a) ∀\,s ∈S. (14) We assume that • 0<δmin≤δ≤δmax<∞0< _ ≤δ≤ _ <∞. • |r|≤max|rmin|,|rmax|<∞|r|≤ \|r_ |,|r_ |\<∞. • S and A are finite. Under these assumptions, the long-run reward rate (gain per unit time) of π is well defined and independent of the initial state: ρ(π):=limT→∞π[∑t=1Trt]π[∑t=1Tδt]=π[r]π[δ],π[δ]:=∑sdπ(s)δ¯π(s),ρ(π)\;:=\; _T→∞ E_π [ _t=1^Tr_t ]E_π [ _t=1^T _t ]\;=\; E_π[r]E_π[δ], _π[δ]\;:=\; _sd^π(s)\, δ^π(s), (15) with π[r]:=∑sdπ(s)∑aπ(a∣s)r¯(s,a)E_π[r]:= _sd^π(s) _aπ(a s) r(s,a). The differential state- and action-value functions of π are Vπ(s)=π[r−ρ(π)δ+Vπ(s′)∣s],Qπ(s,a)=[r−ρ(π)δ+Vπ(s′)∣s,a],V^π(s)=E_π [\,r-ρ(π)δ+V^π(s ) s\, ], Q^π(s,a)=E [\,r-ρ(π)δ+V^π(s ) s,a\, ], (16) so that Vπ(s)=∑aπ(a∣s)Qπ(s,a)V^π(s)= _aπ(a s)Q^π(s,a). Proof for the policy gradient. ∇θρ(πθ)=∑sdπθ(s)∑a∇θπθ(a∣s)Qπθ(s,a)πθ[δ], _θ\,ρ( _θ)= _sd _θ(s) _a _θ _θ(a s)Q _θ(s,a)E_ _θ[δ], (17) where Vπ(s)=π[r−ρ(π)δ+Vπ(s′)∣s],Qπ(s,a)=π[r−ρ(π)δ+Vπ(s′)∣s,a].V^π(s)=E_π [r-ρ(π)δ+V^π(s ) s ], Q^π(s,a)=E_π [r-ρ(π)δ+V^π(s ) s,a ]. Proof. For simplicity, we drop θ when clear from context; all gradients are with respect to θ. Differentiating Vπ(s)=∑aπ(a∣s)Qπ(s,a)V^π(s)= _aπ(a s)Q^π(s,a) and applying the product rule, ∇Vπ(s) ∇ V^π(s) =∑a[∇π(a∣s)Qπ(s,a)+π(a∣s)∇Qπ(s,a)]. = _a [∇π(a s)\,Q^π(s,a)+π(a s)\,∇ Q^π(s,a) ]. (18) By definition, Qπ(s,a)=r¯(s,a)−ρ(π)δ¯(s,a)+∑s′p(s′∣s,a)Vπ(s′)Q^π(s,a)= r(s,a)-ρ(π) δ(s,a)+ _s p(s s,a)V^π(s ). Since the transition kernel does not depend on θ, we have ∇r¯(s,a)=0∇ r(s,a)=0 and ∇δ¯(s,a)=0∇ δ(s,a)=0, hence ∇Qπ(s,a)=−∇ρ(π)δ¯(s,a)+∑s′p(s′∣s,a)∇Vπ(s′).∇ Q^π(s,a)\;=\;-\,∇ρ(π)\, δ(s,a)\;+\; _s p(s s,a)\,∇ V^π(s ). (19) Substituting Equation 19 into Equation 18 and using Equation 13, ∇Vπ(s)=∑a[∇π(a∣s)Qπ(s,a)+π(a∣s)∑s′p(s′∣s,a)∇Vπ(s′)]−∇ρ(π)δ¯π(s).∇ V^π(s)= _a [∇π(a s)Q^π(s,a)+π(a s) _s p(s s,a)∇ V^π(s ) ]\;-\;∇ρ(π)\, δ^π(s). (20) Multiplying Equation 20 by dπ(s)d^π(s) and summing over s∈s∈S, ∑sdπ(s)∇Vπ(s) _sd^π(s)∇ V^π(s) =∑sdπ(s)∑a∇π(a∣s)Qπ(s,a) = _sd^π(s) _a∇π(a s)Q^π(s,a) +∑sdπ(s)∑aπ(a∣s)∑s′p(s′∣s,a)∇Vπ(s′)⏟(⋆)−∇ρ(π)∑sdπ(s)δ¯π(s). + _sd^π(s) _aπ(a s) _s p(s s,a)∇ V^π(s )_( )\;-\;∇ρ(π) _sd^π(s)\, δ^π(s). (21) Exchanging the order of summation in (⋆)( ) and using Equation 14 gives, (⋆)=∑s′[∑sdπ(s)∑aπ(a∣s)p(s′∣s,a)]∇Vπ(s′)=∑s′dπ(s′)∇Vπ(s′).( )\;=\; _s [ _sd^π(s) _aπ(a s)p(s s,a) ]∇ V^π(s )\;=\; _s d^π(s )∇ V^π(s ). (22) The left-hand side of Equation 21 and the term Equation 22 are the same finite quantity, they therefore cancel, leaving 0=∑sdπ(s)∑a∇π(a∣s)Qπ(s,a)−∇ρ(π)π[δ],0\;=\; _sd^π(s) _a∇π(a s)Q^π(s,a)\;-\;∇ρ(π)\,E_π[δ], (23) where we used ∑sdπ(s)δ¯π(s)=π[δ] _sd^π(s) δ^π(s)=E_π[δ] from Equation 15. Since π[δ]≥δmin>0E_π[δ]≥ _ >0 by the boundeness assumption, we may divide by it, which completes the proof. ∎ A.2 More on György et al. (2007) We give the explicit form of the two confidence terms bt(u)b_t(u) and ct(a,s)c_t(a,s) appearing in Equations 7 and 8 (György et al., 2007). Let Tu(t)=∑j=1taj=u(sj),T_u(t)= _j=1^t 1\a_j=u(s_j)\, (24) be the number of rounds on which the played action agreed with policy u, and Ta(s,t)=∑j=1taj=a,sj=s,T_a(s,t)= _j=1^t 1\a_j=a,\,s_j=s\, (25) be the number of times action a was played in context s. Both widths are built from the common scale factor ωt,n=log(t||+1)n,||=||||. _t,n\;=\; (t |U|+1 )n, |U|=|A|^|S|. (26) The width on the rate estimate is bt(u)=ωt,Tu(t)2κ,κ= 2max(rmax−rmin)2δmin2,rmax2(δmax−δmin)2δmin4,b_t(u)\;= _t,\,T_u(t)\; 2κ, κ\;=\;2 \ (r_ -r_ )^2 _ ^2,\;\; r_ ^2\,( _ - _ )^2 _ ^4 \, (27) and the width on the action index is ct(a,s)=(α0+α1)ωt,Ta(s,t),α1=δmax2κ,c_t(a,s)\;=\;( _0+ _1)\, _t,\,T_a(s,t), _1= _ 2κ, (28) α0=8max(rmax−rmin)2,rmax2(δmax−δmin)2δmin2. _0= 8 \(r_ -r_ )^2,\;\; r_ ^2\,( _ - _ )^2 _ ^2 \. (29) The two widths are indexed differently: bt(u)b_t(u) shrinks in the number of rounds consistent with policy u, while ct(a,s)c_t(a,s) shrinks in the number of pulls of action a in context s. The former is subtracted and the latter added, since ρ^t ρ_t enters Equation 8 with a negative sign; pessimism about ρ∗ρ^* is optimism about q∗q^*. Appendix B Theoretical Analysis B.1 Preliminaries Setup and notation. We use the same notation as other parts of our paper, which are discussed in Section 2. Our setting is same as the bandit setting in Section 2.2. Reward rate and value. A stochastic policy π has reward rate ρ(π):=p[∑aπ(a∣x)r¯(x,a)]p[∑aπ(a∣x)δ¯(x,a)],ρ∗:=maxu∈ρ(u),ρ(π)\;:=\; E_p [ _aπ(a x)\, r(x,a) ]E_p [ _aπ(a x)\, δ(x,a) ], ρ^*\;:=\; _u∈Uρ(u), where the maximum runs over the deterministic policies =u:→U=\u:X→A\ of Equation 5. For ρ∈ℝρ we define the relative values qρ(x,a):=r¯(x,a)−ρδ¯(x,a),Mρ(x):=maxaqρ(x,a),F(ρ):=p[Mρ(x)],q_ρ(x,a)\;:=\; r(x,a)-ρ\, δ(x,a), M_ρ(x)\;:=\; _aq_ρ(x,a), F(ρ)\;:=\;E_p [M_ρ(x) ], with q∗:=qρ∗q^*:=q_ρ^* and a∗(x):=argmaxaq∗(x,a)a^*(x):= _aq^*(x,a). Constants. Since rmin≥0r_ ≥ 0, every policy’s reward rate is in [0,ρmax][0, _ ], where ρmax:=rmax/δmin, _ :=r_ / _ , we define the contraction constant as χ:=1−δmin/δmax∈[0,1),χ:=1- _ / _ ∈[0,1), the time spread22 2 We assume the expected time for every (x,a)(x,a) pair cannot be the same since it would reduce the problem into reward maximization in contextual bandit, thus in the SMDP setting we have δspr≠0 _spr≠ 0. We do not count it towards our list of assumptions since it is an intrinsic part of the problem formulation. is δspr:=maxx[maxaδ¯(x,a)−minaδ¯(x,a)]∈(0,δmax−δmin], _spr:= _x [ _a δ(x,a)- _a δ(x,a) ]∈(0, _ - _ ], for a rate ρ∈ℝρ , a greedy policy is denoted as uρ(x)∈argmaxaqρ(x,a),u_ρ(x)∈ _aq_ρ(x,a), thus uρ(x)u_ρ(x) is a greedy action given reward rate ρ in context x, under any fixed tie-breaking. Similarly, we define the greedy margin at ρ as Δ(ρ):=minx[maxaqρ(x,a)−maxa≠uρ(x)qρ(x,a)], (ρ)\;:=\; _x [ _aq_ρ(x,a)- _a≠ u_ρ(x)q_ρ(x,a) ], it is clear that Δ(ρ)≥0 (ρ)≥ 0 generally, and Δ(ρ)>0 (ρ)>0 if and only if the greedy action is unique at every context x. With Δ(ρ∗) (ρ^*), we have the gap at the optimal rate Δ∗:=Δ(ρ∗)=minx[q∗(x,a∗(x))−maxa≠a∗(x)q∗(x,a)]. ^*\;:=\; (ρ^*)\;=\; _x [q^*(x,a^*(x))- _a≠ a^*(x)q^*(x,a) ]. We define the error parameter as ε¯:=Δ∗/(2δspr), := ^*/(2 _spr), since δspr<δmax _spr< _ , we have ε¯>Δ∗/(2δmax). > ^*/(2 _ ). Algorithm. The algorithm is softmax policy with NPG, which keeps a conditional distribution on the actions per context. See Section C.1.1 for more details. Assumption 1 (A1). Reward and time are bounded with r∈[rmin,rmax]r∈[r_ ,r_ ], δ∈[δmin,δmax]δ∈[ _ , _ ], rmin≥0r_ ≥ 0, δmin>0 _ >0. Assumption 2 (A2). Δ(ρ∗)=Δ∗>0 (ρ^*)= ^*>0, which is the positive-gap condition (Auer et al., 2002) for the relative values. Equivalently, for every x, the maximizer a∗(x)a^*(x) of q∗(x,⋅)q^*(x,·) is unique and the gap between the best and second-best relative value is at least Δ∗ ^* at every context. Assumption 3 (A3). Contexts are i.i.d. from p(x)p(x), as described in Section 2.2 and same as in György et al. (2007). Scope. Our analysis focuses on the algorithm’s behavior in expectation, i.e., Q^x,a(ρ)=qρ(x,a) Q_x,a(ρ)=q_ρ(x,a). B.2 Supporting Lemmas Lemma 1. For ρ<ρ′ρ<ρ , δmin≤(F(ρ)−F(ρ′))/(ρ′−ρ)≤δmax; _ \;≤\; (F(ρ)-F(ρ ) )/(ρ -ρ)\;≤\; _ ; in particular F is strictly decreasing. Proof. Let ρ<ρ′ρ<ρ , we have qρ(x,a)−qρ′(x,a)=(r¯(x,a)−ρδ¯(x,a))−(r¯(x,a)−ρ′δ¯(x,a))=(ρ′−ρ)δ¯(x,a),q_ρ(x,a)-q_ρ (x,a)= ( r(x,a)-ρ\, δ(x,a) )- ( r(x,a)-ρ \, δ(x,a) )=(ρ -ρ)\, δ(x,a), let aρ=argmaxaqρ(x,a)a_ρ= _aq_ρ(x,a), aρ′=argmaxaqρ′(x,a)a_ρ = _aq_ρ (x,a), then Mρ(x)=qρ(x,aρ)=qρ′(x,aρ)+(ρ′−ρ)δ¯(x,aρ)≤Mρ′(x)+(ρ′−ρ)δmax,M_ρ(x)=q_ρ(x,a_ρ)=q_ρ (x,a_ρ)+(ρ -ρ) δ(x,a_ρ)≤ M_ρ (x)+(ρ -ρ) _ , similarly Mρ(x)≥qρ(x,aρ′)=qρ′(x,aρ′)+(ρ′−ρ)δ¯(x,aρ′)≥Mρ′(x)+(ρ′−ρ)δmin.M_ρ(x)≥ q_ρ(x,a_ρ )=q_ρ (x,a_ρ )+(ρ -ρ) δ(x,a_ρ )≥ M_ρ (x)+(ρ -ρ) _ . Rearrange, and average over every context x then gives (ρ′−ρ)δmin≤F(ρ)−F(ρ′)≤(ρ′−ρ)δmax.(ρ -ρ) _ ≤ F(ρ)-F(ρ )≤(ρ -ρ) _ . ∎ Lemma 2. F(ρ∗)=0F(ρ^*)=0, so ρ∗ρ^* is the unique root of F; ρ(π)≤ρ∗ρ(π)≤ρ^* for every stochastic policy π. Proof. F(ρ∗)=0F(ρ^*)=0 can be established by Bellman optimality equation, and F is strictly decreasing, so ρ∗ρ^* is unique. For arbitrary policy π, we have (ρ(π)−ρ∗)D(π)=p[∑aπ(a∣x)q∗(x,a)]≤p[∑aπ(a∣x)Mρ∗(x)]=F(ρ∗)=0, (ρ(π)-ρ^* )D(π)=E_p [ _aπ(a x)q^*(x,a) ] _p [ _aπ(a x)M_ρ^*(x) ]=F(ρ^*)=0, where D(π)=p[∑aπ(a∣x)δ¯(x,a)]D(π)=E_p [ _aπ(a x) δ(x,a) ], thus ρ(π)≤ρ∗ρ(π)≤ρ^*. ∎ Lemma 3. Let uρ(x)∈argmaxaqρ(x,a)u_ρ(x)∈ _aq_ρ(x,a). If ρ≤ρ∗ρ≤ρ^*, then ρ≤ρ(uρ)≤ρ∗ρ≤ρ(u_ρ)≤ρ^* and ρ∗−ρ(uρ)≤χ(ρ∗−ρ).ρ^*-ρ(u_ρ)\;≤\;χ\,(ρ^*-ρ). Proof. Let uρ(x)=argmaxaqρ(x,a)u_ρ(x)= _aq_ρ(x,a) be the greedy policy. We have F(ρ)=(ρ(uρ)−ρ)D(uρ),ρ(uρ)=ρ+F(ρ)D(uρ),F(ρ)= (ρ(u_ρ)-ρ )D(u_ρ), ρ(u_ρ)=ρ+ F(ρ)D(u_ρ), by Lemma 1, and that F(ρ∗)=0F(ρ^*)=0, we have δmin(ρ∗−ρ)≤F(ρ)−F(ρ∗)=F(ρ)≤δmax(ρ∗−ρ), _ (ρ^*-ρ)≤ F(ρ)-F(ρ^*)=F(ρ)≤ _ (ρ^*-ρ), since ρ≤ρ∗ρ≤ρ^*, F(ρ)≥δmin(ρ∗−ρ)≥0F(ρ)≥ _ (ρ^*-ρ)≥ 0, thus ρ(uρ)≥ρ(u_ρ)≥ρ, and by Lemma 2, ρ≤ρ(uρ)≤ρ∗ρ≤ρ(u_ρ)≤ρ^*. For the contraction, we have ρ∗−ρ(uρ) ρ^*-ρ(u_ρ) =(ρ∗−ρ)−F(ρ)D(uρ) =(ρ^*-ρ)- F(ρ)D(u_ρ) ≤(ρ∗−ρ)−F(ρ)δmax ≤(ρ^*-ρ)- F(ρ) _ ≤(ρ∗−ρ)−δmin(ρ∗−ρ)δmax ≤(ρ^*-ρ)- _ \,(ρ^*-ρ) _ =(1−δminδmax)(ρ∗−ρ)=χ(ρ∗−ρ). = (1- _ _ )(ρ^*-ρ)=χ\,(ρ^*-ρ). ∎ Lemma 4. For all ρ,ρ′∈ℝρ,ρ , every x, and every a≠uρ′(x)a≠ u_ρ (x), qρ(x,uρ′(x))−qρ(x,a)≥Δ(ρ′)−|ρ−ρ′|δspr.q_ρ (x,u_ρ (x) )-q_ρ(x,a)≥ (ρ )-|ρ-ρ |\, _spr. Consequently Δ(ρ)≥Δ(ρ′)−|ρ−ρ′|δspr (ρ)≥ (ρ )-|ρ-ρ |\, _spr and hence, by symmetry, |Δ(ρ)−Δ(ρ′)|≤|ρ−ρ′|δspr| (ρ)- (ρ )|≤|ρ-ρ |\, _spr: the greedy margin is δspr _spr-Lipschitz in the rate. If moreover |ρ−ρ′|δspr<Δ(ρ′)|ρ-ρ |\, _spr< (ρ ), then uρ=uρ′u_ρ=u_ρ point-wise, under any tie-breaking. Setting ρ′=ρ∗ρ =ρ^* gives, Δ(ρ)≥Δ∗−|ρ−ρ∗|δspr,anduρ=a∗ whenever |ρ−ρ∗|δspr<Δ∗, (ρ)≥ ^*-|ρ-ρ^*|\, _spr, u_ρ=a^*\ whenever \ |ρ-ρ^*|\, _spr< ^*, in particular |ρ−ρ∗|≤ε¯|ρ-ρ^*|≤ gives uρ=a∗u_ρ=a^* and Δ(ρ)≥Δ∗/2 (ρ)≥ ^*/2. Proof. Compare the relative values at two reward-rates ρ,ρ′ρ,ρ similar to Lemma 1 gives qρ(x,a)=qρ′(x,a)+(ρ′−ρ)δ¯(x,a),∀x∈,a∈,q_ρ(x,a)=q_ρ (x,a)+(ρ -ρ) δ(x,a), ∀ x∈X,a∈A, denote u:=uρ′(x)u:=u_ρ (x) as the greedy policy, and let a≠ua≠ u, and substitute the equation above gives qρ(x,u)−qρ(x,a)=[qρ′(x,u)−qρ′(x,a)]⏟(1)+(ρ′−ρ)[δ¯(x,u)−δ¯(x,a)]⏟(2).q_ρ(x,u)-q_ρ(x,a)= [q_ρ (x,u)-q_ρ (x,a) ]_(1)+ (ρ -ρ) [ δ(x,u)- δ(x,a) ]_(2). By definition (see Section B.1), (1)≥Δ(ρ′)(1)≥ (ρ ). For (2)(2), both δ¯ δ terms lie within interval [minbδ¯(x,b),maxbδ¯(x,b)][ _b δ(x,b), _b δ(x,b)], whose length is at most δspr _spr (Section B.1), this gives |(ρ′−ρ)[δ¯(x,u)−δ¯(x,a)]|≤|ρ−ρ′|[maxbδ¯(x,b)−minbδ¯(x,b)]≤|ρ−ρ′|δspr, |(ρ -ρ) [ δ(x,u)- δ(x,a) ] |≤|ρ-ρ |\, [ _b δ(x,b)- _b δ(x,b) ]≤|ρ-ρ |\, _spr, which implies qρ(x,u)−qρ(x,a)≥Δ(ρ′)−|ρ−ρ′|δspr.q_ρ(x,u)-q_ρ(x,a)≥ (ρ )-|ρ-ρ |\, _spr. If the right hand side is positive, then it implies that u is an unique maximizer, thus uρ′(x)=uρ(x)u_ρ (x)=u_ρ(x), then taking the minimum over x gives Δ(ρ)≥Δ(ρ′)−|ρ−ρ′|δspr, (ρ)≥ (ρ )-|ρ-ρ |\, _spr, If the right hand side is non-positive, then the above inequality still holds since Δ(ρ)≥0 (ρ)≥ 0 by construction (Section B.1). By symmetry, we also have that Δ(ρ′)≥Δ(ρ)−|ρ−ρ′|δspr (ρ )≥ (ρ)-|ρ-ρ |\, _spr, which implies |Δ(ρ)−Δ(ρ′)|≤|ρ−ρ′|δspr.| (ρ)- (ρ )|≤|ρ-ρ |\, _spr. This completes the proof for δspr _spr-Lipschitz statement. If |ρ−ρ′|δspr<Δ(ρ′)|ρ-ρ | _spr< (ρ ), then qρ(x,u)−qρ(x,a)>0q_ρ(x,u)-q_ρ(x,a)>0, and since u:=uρ′(x)u:=u_ρ (x), and by uniqueness of the greedy action assumption (Assumption 2), we have uρ′(x)=uρ(x)u_ρ (x)=u_ρ(x) for every x, i.e., the greedy policy for the two rates are the same. For the last statement, substitute ρ′=ρ∗ρ =ρ^* gives the property Δ(ρ)≥Δ∗−|ρ−ρ∗|δspr,anduρ=a∗ whenever |ρ−ρ∗|δspr<Δ∗, (ρ)≥ ^*-|ρ-ρ^*|\, _spr, u_ρ=a^*\ whenever \ |ρ-ρ^*|\, _spr< ^*, and with ε¯=Δ∗/(2δspr) = ^*/(2 _spr), if |ρ−ρ∗|≤ε¯|ρ-ρ^*|≤ , then |ρ−ρ∗|δspr≤Δ∗/2<Δ(ρ∗)=Δ∗|ρ-ρ^*| _spr≤ ^*/2< (ρ^*)= ^*, which implies uρ=a∗u_ρ=a^* and Δ(ρ)≥Δ∗/2 (ρ)≥ ^*/2. ∎ Lemmas 1, 2 and 3 are statements on the Dinkelbach’s algorithm for fractional programming similar to those in Jagannathan (1966); Dinkelbach (1967); Schaible (1976). B.3 Main Results Theorem 3. Let γ∈[0,1]γ∈[0,1] and ρ1∈[0,ρmax] _1∈[0, _ ]. For every k≥1k≥ 1, let uρk(x)∈argmaxaqρk(x,a)u_ _k(x)∈ _aq_ _k(x,a) be a greedy policy under any tie-breaking and ρ(uρk)ρ(u_ _k) is the reward rate of the greedy policy, let ρ~k+1 ρ_k+1 be any value with the sequence satisfying the condition ρ(uρk)1+γ≤ρ~k+1≤(1+γ)ρ(uρk),ρk+1=clip[0,ρmax](ρ~k+1), ρ(u_ _k)1+γ≤ ρ_k+1≤(1+γ)\,ρ(u_ _k), _k+1=clip_[0, _ ] ( ρ_k+1 ), (30) where clip[0,ρmax](z):=minmaxz,0,ρmaxclip_[0, _ ](z):= \ \z,0\,\, _ \. Then, for every k≥2k≥ 2, (1−Rγ1+γ)ρ∗−(χ1+γ)k−2ρmax≤ρk≤(1+γ)ρ∗, (1- Rγ1+γ )ρ^*- ( χ1+γ )^k-2 _ ≤ _k≤(1+γ)ρ^*, (31) where R=δmax/δmin≥1R= _ / _ ≥ 1. If γ=0γ=0, then the sequence increases monotonically ρ2≤ρ3≤⋯≤ρ∗ _2≤ _3≤·s≤ρ^* and the iteration reaches ρk=ρ∗ _k=ρ^* after finite steps. Proof. Throughout the proof, we use the notation N(π):=p[∑aπ(a∣x)r¯(x,a)],D(π):=p[∑aπ(a∣x)δ¯(x,a)],N(π):=E_p [ _aπ(a x)\, r(x,a) ], D(π):=E_p [ _aπ(a x)\, δ(x,a) ], and ρ(π)=N(π)/D(π)ρ(π)=N(π)/D(π) by definition, and p[∑aπ(a∣x)qρ(x,a)] _p [ _aπ(a x)\,q_ρ(x,a) ] =N(π)−ρD(π) =N(π)\;-\;ρ\,D(π) =ρ(π)D(π)−ρD(π)=(ρ(π)−ρ)D(π). =ρ(π)\,D(π)\;-\;ρ\,D(π)= (ρ(π)-ρ )\,D(π). By Lemma 1, we have 0≤ρ(uρk)≤ρ∗≤ρmax.0≤ρ(u_ _k)≤ρ^*≤ _ . We first note that the clipclip operation survives the upper and lower bounds. For upper bounds with a nonnegative right hand side: if z≤Bz≤ B and B≥0B≥ 0, we have clip[0,ρmax](z)≤maxz,0≤B,clip_[0, _ ](z)≤ \z,0\≤ B, and for lower bounds with a right-hand side at most ρmax _ : if z≥bz≥ b with b≤ρmaxb≤ _ , then clip[0,ρmax](z)≥minz,ρmax≥minb,ρmax=b.clip_[0, _ ](z)≥ \z,\, _ \≥ \b,\, _ \=b. Since the bounds survive, for notation simplicity, we use ρk+1 _k+1 and ρ~k+1 ρ_k+1 interchangeably. For k≥1k≥ 1, using the condition on the sequence in Equation 30, we have ρ~k+1≤(1+γ)ρ(uρk)≤(1+γ)ρ∗⟹ρk+1≤(1+γ)ρ∗, ρ_k+1≤(1+γ)\,ρ(u_ _k)\;≤\;(1+γ)\,ρ^* _k+1≤(1+γ)ρ^*, and this completes the upper bound in Equation 31. To prove the lower bound in Equation 31, we define ek:=ρ∗−ρk,B:=Rγρ∗1+γ.e_k:=ρ^*- _k, B:= Rγρ^*1+γ. To complete the rest of the proof, we enumerate each of the two cases when k≥2k≥ 2: (i) the below case with ρk≤ρ∗ _k≤ρ^*. (i) the above case with ρ∗<ρk≤(1+γ)ρ∗ρ^*< _k≤(1+γ)ρ^*. For the below case with ρk≤ρ∗ _k≤ρ^*, using Lemma 3 gives ρ(uρk)≥ρ∗−χ(ρ∗−ρk)=ρ∗−χek,ρ(u_ _k)≥ρ^*-χ(ρ^*- _k)=ρ^*-χ e_k, using the condition in Equation 30, we have ρk+1≥ρ(uρk)/(1+γ) _k+1≥ρ(u_ _k)/(1+γ), which gives ρ∗−ρk+1≤ρ∗−ρ∗−χek1+γ=(1+γ)ρ∗−ρ∗+χek1+γ=χek+γρ∗1+γ,ρ^*- _k+1≤ρ^*- ρ^*-χ e_k1+γ= (1+γ)ρ^*-ρ^*+χ e_k1+γ= χ e_k+γρ^*1+γ, thus we established the fact that ek+1≤χek+γρ∗1+γe_k+1≤ χ e_k+γρ^*1+γ for the below case. For the above case, ρ∗<ρk≤(1+γ)ρ∗ρ^*< _k≤(1+γ)ρ^* implies 0<ρk−ρ∗≤γρ∗0< _k-ρ^*≤γρ^*. By Lemma 1 and the fact that ρ∗<ρkρ^*< _k, we have −δmax(ρk−ρ∗)≤F(ρk)≤−δmin(ρk−ρ∗)≤0.- _ ( _k-ρ^*)≤ F( _k)≤- _ ( _k-ρ^*)≤ 0. (32) Combine together with the fact that F(ρk)/D(uρk)≥F(ρk)/δminF( _k)/D(u_ _k)≥ F( _k)/ _ gives F(ρk)D(uρk)≥−δmaxδmin(ρk−ρ∗), F( _k)D(u_ _k)≥- _ _ ( _k-ρ^*), (33) which implies ρ(uρk) ρ(u_ _k) =ρk+F(ρk)D(uρk)≥ρk−δmaxδmin(ρk−ρ∗) = _k+ F( _k)D(u_ _k)≥ _k- _ _ ( _k-ρ^*) =ρ∗+(ρk−ρ∗)−δmaxδmin(ρk−ρ∗) =ρ^*+( _k-ρ^*)- _ _ ( _k-ρ^*) =ρ∗−(δmaxδmin−1)(ρk−ρ∗)≥ρ∗−(R−1)γρ∗, =ρ^*- ( _ _ -1 )( _k-ρ^*)≥ρ^*- (R-1 )γρ^*, where the first equality comes from the proof in Lemma 3. Use the equation above with the fact ρk+1≥ρ(uρk)/(1+γ) _k+1≥ρ(u_ _k)/(1+γ), we have ρ∗−ρk+1 ρ^*- _k+1 ≤ρ∗−ρ(uρk)1+γ ≤ρ^*- ρ(u_ _k)1+γ (34) ≤ρ∗−ρ∗−(R−1)γρ∗1+γ ≤ρ^*- ρ^*-(R-1)γρ^*1+γ =ρ∗(1+γ)−1+(R−1)γ1+γ=Rγρ∗1+γ=B. =ρ^* (1+γ)-1+(R-1)γ1+γ= Rγρ^*1+γ=B. Hence we established the fact that ek+1≤Be_k+1≤ B for the above case. The final proof is to show by induction that the following is true ek≤χ~k−2ρmax+B,k≥2,e_k≤ χ^k-2 _ +B, k≥ 2, (35) where χ~:=χ/(1+γ)∈[0,1) χ:=χ/(1+γ)∈[0,1). For the base case k=2k=2, since ρ2≥0 _2≥ 0, we have ρ∗−ρ2≤ρ∗≤ρmaxρ^*- _2≤ρ^*≤ _ , which gives e2≤ρmax≤ρmax+Be_2≤ _ ≤ _ +B. For the induction step, we assume the claim is true for some k≥2k≥ 2 and enumerate each of the two cases ρk≤ρ∗ _k≤ρ^* and ρk>ρ∗ _k>ρ^*. If ρk>ρ∗ _k>ρ^*, then by Equation 34, we have ek+1≤B≤χ~k−1ρmax+Be_k+1≤ B≤ χ^k-1 _ +B. If ρk≤ρ∗ _k≤ρ^*, then ek+1 e_k+1 ≤χek+γρ∗1+γ≤χ(χ~k−2ρmax+B)+γρ∗1+γ ≤ χ e_k+γρ^*1+γ≤ χ ( χ^k-2 _ +B )+γρ^*1+γ (36) =χ~k−1ρmax+χB+γρ∗1+γ, = χ^k-1 _ + χ B+γρ^*1+γ, χB+γρ∗1+γ=γρ∗(1+γ)2(χR+1+γ)≤γρ∗(1+γ)2(R+Rγ)=B, χ B+γρ^*1+γ= γρ^*(1+γ)^2(χ R+1+γ)≤ γρ^*(1+γ)^2(R+Rγ)=B, (37) where the last inequality comes from the identity that χR+1=Rχ R+1=R, so χR+1+γ≤R+Rγχ R+1+γ≤ R+Rγ. Combine Equation 36 with Equation 37, we have ek+1≤χ~k−1ρmax+χB+γρ∗1+γ≤χ~k−1ρmax+B,e_k+1≤ χ^k-1 _ + χ B+γρ^*1+γ≤ χ^k-1 _ +B, rearranging the terms gives the lower bound and completes the bound ρ∗(1−Rγ1+γ)−(χ1+γ)k−2ρmax≤ρk.ρ^* (1- Rγ1+γ )- ( χ1+γ )^k-2 _ ≤ _k. We now show that when γ=0γ=0, then the iteration reaches ρ∗ρ^* in finite steps. If γ=0γ=0, we have that every rate estimate equals to the current greedy policy’s rate ρk+1=ρ(uρk) _k+1=ρ(u_ _k), by the identity in Equation 33 and the upper bound ρk≤ρ∗ _k≤ρ^* in Equation 31, we have ρk+1−ρk=F(ρk)D(uρk)≥δminδmax(ρ∗−ρk)≥0, _k+1- _k= F( _k)D(u_ _k)≥ _ _ (ρ^*- _k)≥ 0, hence the sequence increases monotonically. To prove that it reaches ρ∗ρ^* within finite steps, suppose for sake of contradiction that ρk\ _k\ does not reach ρ∗ρ^* in finite steps, then we would have ρk<ρ∗ _k<ρ^* for every iteration k=2,⋯,M+2k=2,·s,M+2, which implies ρk+1−ρk=F(ρk)D(uρk)≥δminδmax(ρ∗−ρk)>0, _k+1- _k= F( _k)D(u_ _k)≥ _ _ (ρ^*- _k)>0, which means the sequence ρk\ _k\ is strictly monotonic. Since greedy policy is a deterministic policy, and there are M=||||M=|A|^|X| possible deterministic policies, this means the possible number of reward rate is |ρ(u):u∈|≤M|\ρ(u):u∈U\|≤ M. Then we would have ρ2<ρ3<⋯<ρM+2 _2< _3<·s< _M+2, which are M+1M+1 distinct rates, a contradiction. ∎ Proposition 1. Let ρ¯∈ℝ ρ be any rate with Δ(ρ¯)>0 ( ρ)>0, write u:=uρ¯u:=u_ ρ for the greedy policy at ρ¯ ρ, which is then unique at every context, and run the algorithm from θ(0)θ^(0) with π0:=πθ(0)(u(x)∣x) _0:= _θ^(0)(u(x) x). For every x and k≥0k≥ 0, 1−πθ(k)(u(x)∣x)≤1−π0π0exp(−ηΔ(ρ¯)k),1- _θ^(k)(u(x) x)\;≤\; 1- _0 _0\, (-η\, ( ρ)\,k ), (38) in particular 1−πθ(k)(u(x)∣x)≤(||−1)exp(−ηΔ(ρ¯)k)1- _θ^(k)(u(x) x)≤(|A|-1) (-η\, ( ρ)\,k ) when π0=1/|| _0=1/|A| (policy is initially uniform). Proof. We use same abbreviations as before, where the greedy policy is denoted as u=uρ^(x)u=u_ ρ(x). By definition of the margin, we have qρ¯(x,u)−qρ¯(x,a)≥qρ¯(x,u)−maxb≠uqρ¯(x,b)≥Δ(ρ¯)>0,for every a≠u,q_ ρ(x,u)-q_ ρ(x,a)≥ q_ ρ(x,u)- _b≠ uq_ ρ(x,b)≥ ( ρ)>0, every a≠ u, the algorithm updates the logit for each step by θx(n+1)=θx(n)+η(qρ¯(x,⋅)−((n)⊤qρ¯(x,⋅)) 1),θ^(n+1)_x=θ^(n)_x+η (q_ ρ(x,·)- ( π^(n) q_ ρ(x,·) )\, 1 ), and the logit between actions u and a is θx,u(n+1)−θx,a(n+1) θ^(n+1)_x,u-θ^(n+1)_x,a =(θx,u(n)+ηqρ¯(x,u)−η(n)⊤qρ¯(x,⋅)) = (θ^(n)_x,u+η\,q_ ρ(x,u)-η\, π^(n) q_ ρ(x,·) ) −(θx,a(n)+ηqρ¯(x,a)−η(n)⊤qρ¯(x,⋅)) \ - (θ^(n)_x,a+η\,q_ ρ(x,a)-η\, π^(n) q_ ρ(x,·) ) =(θx,u(n)−θx,a(n))+η(qρ¯(x,u)−qρ¯(x,a)), = (θ^(n)_x,u-θ^(n)_x,a )\;+\;η (q_ ρ(x,u)-q_ ρ(x,a) ), write d(n):=θx,u(n)−θx,a(n)d^(n):=θ^(n)_x,u-θ^(n)_x,a for the logit gap and c:=η(qρ¯(x,u)−qρ¯(x,a))c\;:=\;η (q_ ρ(x,u)-q_ ρ(x,a) ) for the update gap, the parameter update rule thus becomes a recursion of the form d(n+1)=d(n)+cd^(n+1)=d^(n)+c. Writing as a telescoping sum over n=0,…,k−1n=0,…,k-1 gives d(k)−d(0)=∑n=0k−1(d(n+1)−d(n))=∑n=0k−1c=kc,d^(k)-d^(0)= _n=0^k-1 (d^(n+1)-d^(n) )= _n=0^k-1c\;=\;k\,c, thus we have θx,u(k)−θx,a(k)=θx,u(0)−θx,a(0)+ηk(qρ¯(x,u)−qρ¯(x,a)),θ^(k)_x,u-θ^(k)_x,a=θ^(0)_x,u-θ^(0)_x,a+η\,k (q_ ρ(x,u)-q_ ρ(x,a) ), since our distribution is parameterized through softmax, we have πθ(a∣x)πθ(u∣x)=exp(θx,a)/∑bexp(θx,b)exp(θx,u)/∑bexp(θx,b)=exp(θx,a−θx,u), _θ(a x) _θ(u x)= ( _x,a)/ _b ( _x,b) ( _x,u)/ _b ( _x,b)= ( _x,a- _x,u), fix a≠ua≠ u, and apply the two equations together, we have πθ(k)(a∣x)πθ(k)(u∣x) _θ^(k)(a x) _θ^(k)(u x) =exp(θx,a(k)−θx,u(k)) = (θ^(k)_x,a-θ^(k)_x,u ) =exp(θx,a(0)−θx,u(0)−ηk(qρ¯(x,u)−qρ¯(x,a))) = (θ^(0)_x,a-θ^(0)_x,u-η\,k (q_ ρ(x,u)-q_ ρ(x,a) ) ) =πθ(0)(a∣x)πθ(0)(u∣x)exp(−ηk(qρ¯(x,u)−qρ¯(x,a))), = _θ^(0)(a x) _θ^(0)(u x)\; (-η\,k (q_ ρ(x,u)-q_ ρ(x,a) ) ), observe that −ηk(qρ¯(x,u)−qρ¯(x,a))≤−ηkΔ(ρ¯),-η\,k (q_ ρ(x,u)-q_ ρ(x,a) )\;≤\;-η\,k\, ( ρ), thus πθ(k)(a∣x)πθ(k)(u∣x)≤πθ(0)(a∣x)πθ(0)(u∣x)exp(−ηΔ(ρ¯)k), _θ^(k)(a x) _θ^(k)(u x)\;≤\; _θ^(0)(a x) _θ^(0)(u x)\; (-η ( ρ)k), taking the sum over all a≠ua≠ u gives the following for the left hand side ∑a≠uπθ(k)(a∣x)πθ(k)(u∣x)=∑a≠uπθ(k)(a∣x)πθ(k)(u∣x)=1−πθ(k)(u∣x)πθ(k)(u∣x), _a≠ u _θ^(k)(a x) _θ^(k)(u x)= _a≠ u _θ^(k)(a x) _θ^(k)(u x)= 1- _θ^(k)(u x) _θ^(k)(u x), and the following for the right hand side ∑a≠uπθ(0)(a∣x)πθ(0)(u∣x)exp(−ηΔ(ρ¯)k)=1−π0π0exp(−ηΔ(ρ¯)k), _a≠ u _θ^(0)(a x) _θ^(0)(u x)\; (-η ( ρ)k)= 1- _0 _0\; (-η ( ρ)k), together we have 1−πθ(k)(u∣x)πθ(k)(u∣x)≤1−π0π0exp(−ηΔ(ρ¯)k), 1- _θ^(k)(u x) _θ^(k)(u x)\;≤\; 1- _0 _0\; (-η ( ρ)k), multiplying both sides by πθ(k)(u∣x) _θ^(k)(u x) gives 1−πθ(k)(u∣x) 1- _θ^(k)(u x) =πθ(k)(u∣x)1−πθ(k)(u∣x)πθ(k)(u∣x) = _θ^(k)(u x)\; 1- _θ^(k)(u x) _θ^(k)(u x) ≤πθ(k)(u∣x)1−π0π0exp(−ηΔ(ρ¯)k) ≤ _θ^(k)(u x)\; 1- _0 _0\; (-η ( ρ)k) ≤1−π0π0exp(−ηΔ(ρ¯)k), ≤ 1- _0 _0\; (-η ( ρ)k), and hence completes the proof. ∎ Corollary 1. Let ρ¯∈ℝ ρ be any rate with smallness condition |ρ¯−ρ∗|≤ε¯| ρ-ρ^*|≤ , ε¯=Δ∗/(2δspr) = ^*/(2 _spr). Then uρ¯=a∗u_ ρ=a^*, for every x, a≠a∗(x)a≠ a^*(x), k≥0k≥ 0 1−πθ(k)(a∗(x)∣x)≤1−π0π0exp(−ηΔ∗k/2),1- _θ^(k)(a^*(x) x)≤ 1- _0 _0\, (-η\, ^*k/2 ), (39) in particular 1−πθ(k)(a∗(x)∣x)≤(||−1)exp(−ηΔ∗k/2)1- _θ^(k)(a^*(x) x)≤(|A|-1) (-η ^*k/2) at the zero initialization. Proof. Apply Lemma 4 with Proposition 1. ∎ Theorem 4. Let ρ1,ρ2,… _1, _2,… be a sequence of reward rates. Denote θ(k)θ^(k) as the parameter for the softmax logit after k steps from the zero initialization and ρ¯k=∑n≤kρn/k ρ_k= _n≤ k _n/k, then for every x, k≥1k≥ 1, the policy π(k)π^(k) is a Boltzmann policy of qρ¯k(x,a)q_ ρ_k(x,a) with inverse temperature ηkη k, and if ρk+1=ρ(uρ¯k) _k+1=ρ(u_ ρ_k), ρ1∈[0,ρ∗] _1∈[0,ρ^*], then ρ∗−ρ¯k≤(ρ∗−ρ¯1)exp(1−χ)(k+1)−(1−χ),ρ^*- ρ_k\;≤\;(ρ^*- ρ_1)\, (1-χ)\,(k+1)^-(1-χ), (40) thus ρ¯k ρ_k converges to ρ∗ρ^* at O(k−(1−χ))O(k^-(1-χ)). Moreover, for every k≥k0:=⌈(exp(1−χ)(ρ∗−ρ¯1)/ε¯)1/(1−χ)⌉,k\;≥\;k_0\;:=\; ( (1-χ)\,(ρ^*- ρ_1)/ )^1/(1-χ) , (41) the greedy policy at the running average is optimal, uρ¯k=a∗u_ ρ_k=a^* point-wise, and for every x 1−π(k)(a∗(x)∣x)≤(||−1)exp(−ηΔ∗k/2).1-π^(k)(a^*(x) x)\;≤\;(|A|-1)\, (-η\, ^*k/2 ). (42) Proof. Observe that the relative value qρn(x,a)q_ _n(x,a) for action a and context x at step n can be expressed as ∑n≤kqρn(x,a) _n≤ kq_ _n(x,a) =∑n≤k(r¯(x,a)−ρnδ¯(x,a))=kr¯(x,a)−(∑n≤kρn)δ¯(x,a) = _n≤ k ( r(x,a)- _n\, δ(x,a) )\;=\;k\, r(x,a)\;-\; ( _n≤ k _n ) δ(x,a) =k(r¯(x,a)−ρ¯kδ¯(x,a))=kqρ¯k(x,a), =k\, ( r(x,a)- ρ_k\, δ(x,a) )=k\,q_ ρ_k(x,a), By the parameter update rule for NPG with a softmax policy θx′=θx+η(q−⟨,q⟩)θ _x= _x+η(q- π,q 1), fix the context, and expand the terms through the recursion gives θx(k) θ^(k)_x =η∑n≤kqρn(x,⋅)−η(∑n≤ksn)sn:=⟨(n−1),qρn(x,⋅)⟩ =η _n≤ kq_ _n(x,·)\;-\;η ( _n≤ ks_n ) 1 s_n:= π^(n-1),q_ _n(x,·) =ηkqρ¯k(x,⋅)+ck(x) 1ck(x):=−η∑n≤ksn, =η\,k\,q_ ρ_k(x,·)\;+\;c_k(x)\, 1 c_k(x):=-η _n≤ ks_n, when the logit θx(k) _x^(k) is feed into the softmax function, since softmax is invariant to addition of constants, the ck(x)c_k(x) term is ignored, thus the policy becomes the Boltzmann policy of ηkqρ¯k(x,a)η\,k\,q_ ρ_k(x,a) and completes the first part of the proof. For the second part, first define the estimated rate as kρ¯k=∑n≤kρnk ρ_k= _n≤ k _n, then ρ¯k+1=1k+1∑n≤k+1ρn=kρ¯k+ρk+1k+1=k+1ρ¯k+1k+1ρk+1, ρ_k+1= 1k+1 _n≤ k+1 _n= k\, ρ_k+ _k+1k+1= kk+1\, ρ_k\;+\; 1k+1\, _k+1, which is ρ¯k+1=(1−αk)ρ¯k+αkρk+1,αk=1/(k+1). ρ_k+1=(1- _k) ρ_k+ _k\, _k+1, _k=1/(k+1). (43) Since the algorithm only accepts samples selected by argmaxaθx,a _a _x,a, we assume that the rate estimate equals to the greedy policy’s rate under ρ¯k ρ_k, i.e., ρk+1=ρ(uρ¯k) _k+1=ρ(u_ ρ_k). Since ρ¯k≤ρ∗ ρ_k≤ρ^* at k=1k=1 by assumption, Lemma 3 gives that ρ(uρ¯k)≥ρ¯kρ(u_ ρ_k)≥ ρ_k and ρ∗−ρ(uρ¯k)≤χ(ρ∗−ρ¯k)ρ^*-ρ(u_ ρ_k)≤χ(ρ^*- ρ_k). ρ¯k+1 ρ_k+1 is a convex combination of ρ¯k ρ_k and ρ(uρ¯k)ρ(u_ ρ_k) by Equation 43 and we have ρ¯k≤ρ¯k+1≤ρ∗ ρ_k≤ ρ_k+1≤ρ^*. Subtracting Equation 43 from ρ∗ρ^* gives ρ∗−ρ¯k+1 ρ^*- ρ_k+1 =ρ∗+αkρ∗−αkρ∗−(1−αk)ρ¯k−αkρk+1 =ρ^*+ _kρ^*- _kρ^*-(1- _k) ρ_k- _k _k+1 =(1−αk)(ρ∗−ρ¯k)+αk(ρ∗−ρk+1) =(1- _k)(ρ^*- ρ_k)+ _k(ρ^*- _k+1) ≤(1−αk)(ρ∗−ρ¯k)+αkχ(ρ∗−ρ¯k)=(1−αk(1−χ))(ρ∗−ρ¯k). ≤(1- _k)(ρ^*- ρ_k)+ _kχ(ρ^*- ρ_k)= (1- _k(1-χ) )(ρ^*- ρ_k). Denote ek=ρ∗−ρ¯ke_k=ρ^*- ρ_k, c:=1−χ=δmin/δmax∈(0,1]c:=1-χ= _ / _ ∈(0,1], and expand the recursion gives ek e_k ≤e1∏m=2k(1−c/m)≤e1exp(−c∑m=2k1/m) ≤ e_1 _m=2^k (1-c/m )≤ e_1 (-c _m=2^k1/m ) ≤e1exp(c−clog(k+1))=e1exp(c)(k+1)−c, ≤ e_1 (c-c (k+1) )=e_1 (c)\,(k+1)^-c, where we used the fact 1−z≤exp(−z)1-z≤ (-z) and ∑m≤k1/m≥log(k+1) _m≤ k1/m≥ (k+1). Since ρk+1=ρ(uρ¯k)≤ρ∗ _k+1=ρ(u_ ρ_k)≤ρ^*, ρ¯k+1≤(1−αk)ρ¯k+αkρ∗<ρ∗ ρ_k+1≤(1- _k) ρ_k+ _kρ^*<ρ^* whenever ρ¯k<ρ∗ ρ_k<ρ^*, so by induction ρ1<ρ∗ _1<ρ^* keeps the strict inequality ρ¯k<ρ∗ ρ_k<ρ^* to every k>1k>1. For the last part, the second part gives ρ¯k≤ρ¯k+1≤ρ∗ ρ_k≤ ρ_k+1≤ρ^*, so the error ρ∗−ρ¯kρ^*- ρ_k is non-increasing in k and it suffices to find the first index at which it falls below ε¯ . By Equation 40, ρ∗−ρ¯k≤ε¯ρ^*- ρ_k≤ holds when (k+1)1−χ≥exp(1−χ)(ρ∗−ρ¯1)/ε¯,(k+1)^1-χ≥ (1-χ)\,(ρ^*- ρ_1)/ , thus we need k≥(exp(1−χ)(ρ∗−ρ¯1)/ε¯)1/(1−χ)−1,k≥ ( (1-χ)\,(ρ^*- ρ_1)/ )^1/(1-χ)-1, thus we can set k0k_0 to be the constant stated in Equation 41, and it then holds at every subsequent step k≥k0k≥ k_0. When k≥k0k≥ k_0, since |ρ¯k−ρ∗|≤ε¯| ρ_k-ρ^*|≤ , we can apply Lemma 4, which gives uρ¯k=uρ∗=a∗u_ ρ_k=u_ρ^*=a^* and Δ(ρ¯k)≥Δ∗/2 ( ρ_k)≥ ^*/2. Similar to the steps in the proofs for Propositions 1 and 1, the probability ratio is π(k)(a∣x)π(k)(a∗(x)∣x) π^(k)(a x)π^(k)(a^*(x) x) =exp(−ηk(qρ¯k(x,a∗(x))−qρ¯k(x,a))) = (-η\,k (q_ ρ_k(x,a^*(x))-q_ ρ_k(x,a) ) ) ≤exp(−ηkΔ(ρ¯k)) ≤ (-η\,k\, ( ρ_k) ) ≤exp(−ηΔ∗k/2), ≤ (-η\, ^*\,k/2 ), sum over all actions a≠a∗(x)a≠ a^*(x) gives 1−π(k)(a∗(x)∣x)π(k)(a∗(x)∣x)≤(||−1)exp(−ηΔ∗k/2), 1-π^(k)(a^*(x) x)π^(k)(a^*(x) x)≤(|A|-1) (-η\, ^*\,k/2 ), which gives Equation 42 after multiplying both sides by π(k)(a∗(x)∣x)≤1π^(k)(a^*(x) x)≤ 1 and completes the proof. ∎ Appendix C Implementation C.1 Bandit Experiments This section gives the full specification of the bandit experiments in Section 3.3. C.1.1 Bandit Baselines C-UCB with theoretical constants. The algorithm of György et al. (2007), implemented faithfully for all the constants. C-UCB with tuned constants. Identical in every respect, except that the two scale constants of Equations 27 and 28 are replaced by free parameters c1c_1 and c2c_2 selected via grid search. C-UCB with approximation. C-UCB involves enumeration over all policies for reward-rate estimation. This is infeasible when we have many arms and contexts. As an approximation, we draw 2048 policies uniformly random from policies constructed from visited arms and contexts, and enumerate on these 2048 policies. NPG. The reward-rate learner keeps one independent softmax head per context: a tabular parameter matrix θ∈ℝ||×(K+1)θ ^|X|×(K+1) with one free logit per (x,a)(x,a) pair, initialized to zero so that πθ(a∣x)=1/(K+1) _θ(a x)=1/(K+1) is uniform, and πθ(a∣x)=exp(θx,a)∑bexp(θx,b). _θ(a x)\;=\; ( _x,a) _b ( _x,b). (44) At each round the learner observes x, samples a∼πθ(⋅∣x)a _θ(· x), and receives (r,δ)(r,δ), and its regret is incurred by its own action choices. The update rule is the natural gradient of the relative reward θx←θx+η,∇θxJ(θ)=(θx), _x← _x+η\, z, _ _x\,J(θ)=F( _x) z, (45) where the Fisher information matrix F is (θx)=diag()−⊤,F( _x)=diag( π)- π π , (46) which simplifies to adding the advantage to the logits directly. Writing =πθ(⋅∣x) π= _θ(· x), s for the vector with entries sa=[r−ρ∗δ∣x,a]s_a=E[\,r-ρ^*δ x,a\,], J=⊤J= π^\! s and =−(⊤) a= s- ( π^\! s ) 1 be the advantages ∂πa∂θc ∂ _a∂ _c =πa(a=c−πc)⟹∇θxJ=⊙, = _a ( 1\a=c\- _c ) _ _xJ= π a, (47) solving for =⊙F z= π a gives = z= a as one of the solutions, thus the update rule is θx←θx+ηAθ(x,⋅). _x← _x+η\,A_θ(x,·). The advantage is needed for every arm while a bandit learner observes only the arm it played, so we keep a per-cell buffer of the raw sums Sx,arS^r_x,a and Sx,aδS^δ_x,a with counts nx,an_x,a, and use them to estimate the expected relative reward under current estimated rate ρ^t ρ_t, Q^x,a=(Sx,ar−ρ^tSx,aδ)/nx,a,nx,a>0,0,nx,a=0,θx←θx+η(Q^x−(⊤Q^x)). Q_x,a\;=\; cases (S^r_x,a- ρ_t\,S^δ_x,a )/n_x,a,&n_x,a>0,\\[2.0pt] 0,&n_x,a=0, cases _x\;←\; _x+η ( Q_x-( π^\! Q_x) 1 ). (48) The rate ρ^t ρ_t is supplied by the online NIW estimator. Vanilla policy gradient. Everything is kept same as NPG, except the gradient update is θx←θx+η∇θxJ(θ), _x← _x+η _ _xJ(θ), which gives the update rule θx←θx+η⊙(Q^x−(⊤Q^x)). _x← _x+η\; π ( Q_x-( π Q_x) 1 ). The NIW rate estimator. The estimate ρ^t ρ_t is fitted online on pairs of accumulated reward and accumulated time, mirroring the way the rate is formed from per-slot accumulators in MLE-Bench and NanoGPT (Section 4.1). It is refitted every 88 steps, and ρ^t ρ_t is the 95th percentile among posterior-predictive samples (see Section C.2). The accumulators are gated, in that a sample enters an accumulator only when the arm played was the current argmaxaθa,x _a\, _a,x, so that the estimate describes the greedy policy. They are also age-bounded, in that a slot restarts once it is older than A steps, with A tuned with hyperparameter search. We note that this estimator does not guarantee exact estimation of the greedy policy’s reward rate. Rather, it serves as a test case to see how our algorithm behaves when the greedy reward rate has approximation error, which is the case in empirical settings. C.1.2 Experiment Setup Setup. At each decision epoch the learner selects an action a and observes a reward r≥0r≥ 0 together with a time δ>0δ>0. The pair is drawn jointly, so the reward and the time it costs may be dependent. Performance is measured against the optimal reward rate ρ∗ρ^*, and we report the realized regret of Equation 6 at step t. Each context and action pair consists of a mean reward and time parameter, forming a table of size (K,||)(K,|X|). The setup for these mean parameters are described below. To ensure we do not introduce bias in terms of ordering of the arms and contexts, we perform a permutation on K arm indices and |||X| context indices for every trial. Arm design. The K+1K+1 actions available in a context are the skip action and two structurally different kinds of arm, which are specified in opposite directions. The four designed arms fix the pair ([r(a)],[δ(a)])(E[r(a)],E[δ(a)]) and let the rate follow: • best arm (0.50,1.0)(0.50,1.0) at rate 0.5000.500. • decoy (0.57,1.2)(0.57,1.2) at rate 0.4750.475, carrying higher raw reward at a 5%5\% relative rate deficit. • trap (0.90,3.0)(0.90,3.0) at rate 0.3000.300, the highest raw reward in the context at the worst rate. • fast arm (0.21,0.6)(0.21,0.6) at rate 0.3500.350 that is cheap but weak. The decoy and the trap both exist to punish an agent that maximizes reward rather than reward per unit of time. The remaining K−4K-4 filler arms fix the rate and derive the reward and time from it ζj=βx(0.42−jK−5(0.42−0.18)),[δ(4+j)]=dj,dj∼Unif(0.70, 2.00),[r(4+j)]=ζjdj, gathered _j= _x (0.42- jK-5 (0.42-0.18 ) ), [δ(4+j)]=d_j,\\ d_j (0.70,\,2.00), [r(4+j)]= _j\,d_j, gathered (49) for j=0,…,K−5j=0,…,K-5, where βx _x is the reward scale of the context (see below) and each djd_j is drawn independently. This means the reward rate is monotone but reward and time are unordered, so neither quantity on its own identifies a good arm. No filler arm is ever optimal, because the sequence’s maximum is at 0.42βx0.42\, _x, strictly below the best arm at 0.50βx0.50\, _x. Filler arms therefore widen the action set at controlled rate spacing rather than competing for the optimum, which makes K an exploration-difficulty axis. Context Design. A context x∈x∈X is drawn uniformly at each epoch and adds a skip action with (0,0.35)(0,0.35). A context is summarized by the reward scale βx _x, which multiplies every reward and leaves time untouched; since the rate is reward divided by time, βx _x scales the rate as well. Four contexts are designed: • marginal (βx=1.00 _x=1.00) reproduces the arm design above and serves as the baseline. • rich (βx=1.50 _x=1.50) holds the highest multiplier. • poor (βx=0.32 _x=0.32) so that skip can be considered. • coupling is not a rescaling at all. Instead, it replaces the four designed arms with (1.80,2.00)(1.80,2.00), (1.00,1.00)(1.00,1.00), (1.60,3.0)(1.60,3.0) and (0.20,0.6)(0.20,0.6), a slow-but-rich arm at rate 0.900.90 against a fast-but-lean arm at rate 1.001.00. Maximizing the rate within coupling takes the fast arm while the optimal policy takes the slow one, which is why the problem does not decompose into |||X| independent bandits. This property holds as long as ρ∗<0.8ρ^*<0.8, which is the case for all of our simulations. Any context beyond these four is filler, where βx _x’s are uniformly binned midpoints between [0.42,1.35][0.42,1.35], the range between poor and rich. Environment families. Four families vary the dependence between reward and time while holding the marginal means fixed. Let μar=[ra]μ^r_a=E[r_a] and μaδ=[δa]μ^δ_a=E[ _a] be the designed means of arm a. We fix σr=0.20 _r=0.20, a time noise fraction ς=0.15 =0.15, a log-scale spread σlog=0.5 _ =0.5, and a time floor δmin=0.5 _ =0.5. Every family draws a pair of standard normals z1,ξ∼(0,1)z_1,ξ N(0,1) and couples them through a Gaussian copula z2=ϱz1+1−ϱ2ξ,corr(z1,z2)=ϱ,z_2\;=\; \,z_1\;+\; 1- ^2\;ξ, (z_1,z_2)= , (50) so that ϱ alone controls the reward–time dependence while both marginals are unchanged. E1 (independent), ϱ=0 =0. Reward and time are rectified normals drawn independently, ra=max(μar+σrz1,0),δa=max(μaδ+ςμaδz2,δmin).r_a= (μ^r_a+ _rz_1,0 ), _a= (μ^δ_a+ \,μ^δ_az_2,\; _ ). (51) E2 (positively coupled), ϱ=+0.8 =+0.8 and E3 (negatively coupled), ϱ=−0.8 =-0.8. Identical marginals to Equation 51, with z2z_2 generated by Equation 50. Under E2 a high-reward draw tends to also be slow, so the empirical rate r/δr/δ is compressed and a learner that ranks arms by raw reward is punished less; under E3 high reward tends to arrive fast, which inflates the apparent spread in rates. E4 (lognormal), ϱ=+0.5 =+0.5 on the log scale. Both coordinates are lognormal with heavy right tails. Note that we take the log on the mean parameter so that the expectation [ra],[δa]E[r_a],E[ _a] still equals to μar,μaδμ^r_a,μ^δ_a: ra=exp(logμar−σlog22+σlogz1),δa=max(exp(logμaδ−σlog22+σlogz2),δmin).r_a= \! ( μ^r_a- _ ^22+ _ z_1 ), _a= \! ( \! ( μ^δ_a- _ ^22+ _ z_2 ),\; _ ). (52) The rectification and the floor δmin _ make the realized marginals differ slightly from (μar,μaδ)(μ^r_a,μ^δ_a), so we do not rely on the design values when computing ρ∗ρ^* (See Equation 58). Every reported mean is instead computed in closed form, using [max(X,c)]=cΦ(a)+m(1−Φ(a))+sϕ(a),a=c−ms,E[ (X,c)]=c\, (a)+m (1- (a) )+s\,φ(a), a= c-ms, (53) [max(Y,c)]=cΦ(b)+m(1−Φ(b−σlog)),b=logc−μσlog,E[ (Y,c)]=c\, (b)+m (1- (b- _ ) ), b= c-μ _ , (54) for a rectified normal X∼(m,s2)X N(m,s^2) and a floored lognormal Y with log-scale σlog _ . ϕ,Φφ, are PDF and CDF of Gaussian, respectively. The means are therefore deterministic and do not require Monte Carlo approximation, and ρ∗ρ^* follows from Equation 58. Ground truth. The general average-reward SMDP optimality equation reads V∗(x)=maxa∈[r(x,a)]−ρ∗[δ(x,a)]+∑x′∈p(x′∣x,a)V∗(x′).V^*(x)= _a \E[r(x,a)]-ρ^*\,E[δ(x,a)]+ _x p(x x,a)\,V^*(x ) \. (55) In the bandit case the next context is drawn independently of the current context and of the action taken, so p(x′∣x,a)=p(x′)p(x x,a)=p(x ) and the continuation term collapses to a constant c:=∑x′p(x′)V∗(x′)c:= _x p(x )V^*(x ) that depends on neither x nor a. Writing q∗(x,a):=[r(x,a)]−ρ∗[δ(x,a)],q^*(x,a):=E[r(x,a)]-ρ^*\,E[δ(x,a)], (56) Equation 55 becomes V∗(x)=maxa∈q∗(x,a)+c.V^*(x)= _a q^*(x,a)+c. (57) Averaging against p(x)p(x) on both sides and cancelling c gives the condition that ρ∗ρ^* is the unique root of F(ρ)=∑xp(x)maxa∈([r(x,a)]−ρ[δ(x,a)])= 0.F(ρ)\;=\; _xp(x)\, _a \, (E[r(x,a)]-ρ\,E[δ(x,a)] )\;=\;0. (58) F is a maximum of affine functions of ρ, hence convex, piecewise linear and strictly decreasing, so the root is unique and we find ρ∗ρ^* by bisection. Empirically, the residual is small |F(ρ∗)|∼10−17|F(ρ^*)| 10^-17. Evaluation. Every instance is run to horizon T=30,000T=30,000 over all four families, in two studies whose arm and context counts differ: • The small-scale study at ||=4|X|=4 with K∈5,10,20,30K∈\5,10,20,30\. • The larger-scale study at ||∈8,16,32,64|X|∈\8,16,32,64\ with K∈70,100,130,160K∈\70,100,130,160\. Hyperparameters are selected by grid search on tuning seeds disjoint from the evaluation seeds, taking the median final regret over the tuning seeds, and the selected configuration is then run on 100100 held-out evaluation seeds. Every step that took place counts as part of the regret calculation, including the burn-in phase (for instance, to have a nonzero value of counts for C-UCB). Confidence bands are pointwise percentile bootstraps of the mean over trials. C.2 Reward-rate Predictor We introduce the reward-rate estimator we use for MLE. For fρf_ρ, we use a NIW prior with joint Gaussian likelihood (Murphy, 2012) fitted from historical samples observed in buffer ℬB. We model the pair as (racc,logδacc)(r_acc, _acc), where raccr_acc and δacc _acc are the accumulated reward and time, and we take the logarithm because we need time to be positive. Writing zi=(racc,i,logδacc,i)z_i=(r_acc,i,\, _acc,i) for the ithi^th of the n samples available at step t, where n is the batch size. The likelihood and its conjugate prior are zi|μ,Σ∼(μ,Σ),(μ,Σ)∼NIW(μt−1,λκt−1,λνt−1,λΨt−1),z_i μ, N(μ, ), (μ, ) ( _t-1,\,λ _t-1,\,λ _t-1,\,λ _t-1 ), (59) where μt−1 _t-1 is the prior mean of μ, Ψt−1 _t-1 is the prior scatter matrix of Σ , and κt−1 _t-1 and νt−1 _t-1 are pseudo-counts. The prior at step t is the previous step’s posterior discounted by a forgetting factor λ∈(0,1]λ∈(0,1], which shrinks those pseudo-counts so that old evidence decays geometrically, and the posterior predictive is a Student-t with νt−1 _t-1 degrees of freedom, where νt=λνt−1+n _t=λ _t-1+n is the posterior count after this step’s n samples are absorbed. Reward rate is estimated by sampling M batches of n draws each from this predictive. Writing the mthm^th batch as (xm,j,ym,j)j=1n\(x_m,j,\,y_m,j)\_j=1^n, each batch yields one rate, and ρ^t+1 ρ_t+1 is the 95th percentile of them, ρm=∑j=1nxm,j∑j=1nexp(ym,j),ρ^t+1=Quantile0.95(ρmm=1M). _m= _j=1^nx_m,j _j=1^n (y_m,j), ρ_t+1=Quantile_0.95 (\ _m\_m=1^M ). (60) C.3 MLE-Bench C.3.1 Environment, Reward and Self-improvement Execution. Every action a is a code block generated by the policy (Qwen3.5-4B). The code block is executed in its own sandbox subprocess environment with its own working directory, under a memory budget of 16GB and a compute budget of 2 CPUs. The subprocess is killed when it reaches timeout δmax _ . We set δmax=600 _ =600 seconds, except for nomad2018, pizza, tabular-2021 and tabular-2022, where the dataset is large enough that 600600 seconds admits almost no honest attempt, so we use δmax=1200 _ =1200 seconds for these tasks. Time δ is floored at clip_time =0.1=0.1 seconds before it is used anywhere, so that no sample can claim an unboundedly large rate. We train with a batch size of 128. Thus, assuming full parallelization at each step, we require 16×128 = 2048GB of memory with 256 CPUs. This amount of resource is large in academic setting, so we use a first-in-first-out queue that processes the 128 samples by 16 samples running parallel at a time through Ray. Reward. A run is a valid submission when the script terminates before δmax _ , generates a submission file, and that file is accepted by the MLE-Bench grader. The grader returns the competition’s own metric. For the 1515 tasks whose metric is higher-is-better and bounded, we use performance metric (such as AUC) directly as the reward. For the remaining 77 tasks, where the metric is lower-is-better and has no bounded value (such as log loss), we use the position of the submission on the MLE-Bench leaderboard normalized as percentiles as reward. When there is a tie on the leaderboard, we use the mid-point position. rt+1=1N∑i=1N(ℓi≺g+12 1ℓi=g)∈[0,1],r_t+1\;=\; 1N _i=1^N ( 1\ _i g\\;+\; 12\, 1\ _i=g\ )\;∈\;[0,1], (61) where ℓ1,…,ℓN _1,…, _N are the leaderboard entries and ≺ orders them in the competition’s own direction. This restores the higher-is-better direction and bounds the reward in [0,1][0,1]. The 7 tasks using leaderboard score are exactly the seven marked with (↓)( ) in Table 1. Failures. If a script fails to obtain a valid submission, we define it as a failure. We use a penalty of −10-10 for invalid submissions. Writing s for a valid submission’s score of Equation 61, the reward functions for vanilla and RPG are r~t+1RPG=s−ρ^t+1δt+1,valid,−10−ρ^t+1δmax,invalid,r~t+1vanilla=s,valid,−10,invalid. r^\,RPG_t+1= casess- ρ_t+1\, _t+1,&valid,\\[2.0pt] -10- ρ_t+1\, _ ,&invalid, cases r^\,vanilla_t+1= casess,&valid,\\[2.0pt] -10,&invalid. cases (62) For RPG, we charge ρ^t+1δmax ρ_t+1\, _ for invalid submissions so that every invalid sample’s reward is strict below the valid one. Self-improvement buffer. The transition stored in ℬB is the extracted code, the grader’s message, the score s of the submission, and the measured time δt+1 _t+1. The buffer holds a capacity of 10 times the batch size, i.e. 1,2801,280 solutions. At every step, it is ranked by the relative reward key(a,r,δ)=r−ρ^t+1δ,key(a,r,δ)\;=\;r\;-\; ρ_t+1\,δ, (63) with ρ^t+1 ρ_t+1 obtained from the current step’s reward rate estimate. The vanilla arm uses the identical buffer with ρ^=0 ρ=0. The top 128128 solutions ranked by the key fill the 128128 self-improvement prompts of the next step, cycling from the top when the buffer holds <128<128 samples. A valid sample whose score strictly exceeds every retained solution’s receives an additional self-improvement bonus of 0.50.5, which shapes the advantage of both arms but enters neither Equation 63 nor the rate estimate. C.3.2 Prompts Base prompt. The initial state s0s_0 is a chat-formatted prompt carrying the data path, the task description, the metric, the available packages, the required submission format, and a schema snippet of the training file. It also carries the single efficiency instruction referred to in Section 4.1. See Figure 5 for an example. Example Prompt for spooky-author-identification ⬇ You are given a machine learning task. You must solve the task by training a model and running the model on the test set to produce a submission file. First outline your proposed solution in natural language (3-5 sentences), followed by a single markdown code block (wrapped in ‘) which implements this solution. Note: - Datasets train.csv and test.csv are available in ‘[...]/spooky-author-identification/prepared/public/‘. - You MUST produce a submission file at ‘./submission.csv‘ by running your model on the test split. - To keep runtime and memory low, make your training efficient: train on a subsample of the training rows and/or a reduced set of features (via feature selection or sparse methods) rather than the full data. - Your task is to predict the author (EAP, HPL, MWS) of horror story excerpts. - Submissions are evaluated using multi-class logarithmic loss. - You can use pre-installed packages such as: ‘statsmodels‘, ‘pandas‘, ‘timm‘, ‘bayesian-optimization‘, ‘scikit-learn‘, ‘xgboost‘, ‘numpy‘, ‘torch‘, ‘torchvision‘, ‘lightGBM‘, ‘torch-geometric‘. - submission.csv should have the following format: ‘ id,EAP,HPL,MWS id07943,0.33,0.33,0.33 etc. ‘ - Data snippet: -> [...]/train.csv has 17621 rows and 3 columns. author (object) has 3 unique values: [’EAP’, ’MWS’, ’HPL’] [...] Figure 5: Example prompt on MLE-Bench. Self-improvement prompt. The next state st+1s_t+1 replaces the opening instruction with the self-improvement phrase i0i_0 and appends the selected solution and its grader message. Every other line is byte-identical to the base prompt, so the two states differ only in what the self-improvement loop contributes. The differing lines are shown in Figure 6. previous_plan_code is filled with the executed code of the selected transition, which is the action a. previous_plan_error contains the grader’s own message. For a valid submission, previous_plan_error reports the achieved time δ and performance score s, which is how the agent is told the number it is being asked to beat. For an invalid one, it reports the failure and the traceback. We set the prompt length to 25602560 tokens and response length to 25002500 tokens. Example Prompt for Self-Improvement ⬇ You have implemented a 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 a 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. Note: [... identical to the base prompt ...] Previous solution is ‘previous_plan_code‘ The output or error message for this solution is ‘previous_plan_error‘ Figure 6: Example prompt for self-improvement. previous_plan_code is filled with the executed code of the selected transition, which is the action a. previous_plan_error contains the grader’s own message. C.3.3 Hyperparameters See Table 2. Table 2: MLE-Bench configuration. The upper block is shared by RPG and vanilla. The lower block is where they differ. Group Setting Value Model actor and reference Qwen3.5-4B critic Qwen3.5-4B, token-classification head rollout engine vLLM, bfloat16 attention, checkpointing sdpa, on PPO advantage estimator GAE, γ=1.0γ=1.0, λ=1.0λ=1.0 actor / critic learning rate 1×10−61× 10^-6 / 1×10−51× 10^-5 train batch == mini-batch 128128 micro-batch (actor, critic, reference) 44 clip ratio, value clip 0.20.2, 0.50.5 gradient clip 1.01.0 entropy coefficient 0.0010.001 KL coefficient to the reference 0.0010.001, fixed Rollout temperature, top-p, top-k 0.70.7, 0.80.8, 2020 presence penalty 1.51.5 samples per prompt 11 prompt / response length 25602560 / 25002500 tokens Environment execution timeout δmax _ 600600 s (12001200 s on four tasks) time floor 0.10.1 s memory / CPU per sample 1616 GB / 22 invalid-submission penalty −10-10 Buffer capacity 10×128=1,28010× 128=1,280 solutions selection width top 128128 self-improvement bonus 0.50.5 Compute per run 2×2× A100 80GB, 32 CPU, 450 GB repetitions 33 per task per arm tasks 2222 MLE-Bench Lite competitions RPG only rate estimator online NIW, Equation 59 forgetting factor λ 0.30.3 predictive batches M 1,0001,000 percentile to be selected for M batches 95 time charged to an invalid sample δmax _ C.4 NanoGPT C.4.1 Environment and Grader The task. An action is a modification to a single self-contained train_gpt.py, launched as torchrun --nproc_per_node=8 train_gpt.py on eight A100-SXM4-80GB GPUs and graded against the rules of the modded-NanoGPT speedrun (Jordan et al., 2024). The FineWeb-10B dataset is fixed, and the validation loss must be the mean cross-entropy over exactly 10,485,76010,485,760 validation tokens. A validation pass must be printed every 125125 training steps and again at the last step. A script that alters the validation tokens or the logging interval is scored as a failure, so the agent cannot design hacks. The elapsed time δ is the whole process’s wall clock, measured from launch to exit. The sandbox process is killed at timeout of δmax=1500 _ =1500 seconds. Outcomes. A crash, a parse failure, or a timeout with no usable validation line is an invalid attempt. A run that is killed at timeout but has printed real validation lines is kept as valid and scored on its reached validation loss and charged its measured elapsed time. This is the only case in which the timeout and a valid outcome coexist. Times are floored at 0.10.1 seconds. Seeds. The buffer is initialized with past leaderboard submissions to modded NanoGPT, re-executed on our own eight-A100 grader so that their recorded time is measured on the same clock the agent is scored on. These past leaderboard submissions serve as the roots where every self-improvement chain descends from. Unlike in MLE-Bench, we do not have an initial state s0s_0 since the buffer is non-empty from the first step. C.4.2 Delta Proposals and the Implementer Implementer. A train_gpt.py used for self-improvement prompts is around 560560 lines. Asking the policy (Qwen3.5-4B) to generate such scripts would use a lot of response tokens, slow down training, and the majority of the focus would be on correct code transcription. To alleviate this issue, we ask the policy to propose a change specification in three fixed sections: diagnosis, change, and risk. For each section, it needs to point out the corresponding snippet in the parent code block to be improved upon. The implementer is GPT 5.6 Terra at temperature 00 with a 16,38416,384-token budget, and it is given the parent script, the change specification, and asked for creating a search-and-replace edit blocks. The edit blocks then go through regular expression to replace the parent code with the new edits, which we call the output script as patch. A reply by the implementer is accepted only if every edit block matches the parent exactly and the validation logging lines remain untouched (to prevent model from hacking the validation logs). Judge. A separate judge, the same GPT 5.6 Terra model, gives three votes at temperature 00 to check whether the patch script violates the above constraints. If any of the violations occur, it is scored as invalid. C.4.3 Reward RPG. Writing L for the achieved validation loss and L⋆=3.28L =3.28 for the target, an attempt’s reward is the loss reduction below a fixed reference Lref=10L_ref=10, plus the two threshold bonuses described in Section 4.1, rt+1=max(0,Lref−L)+B+wmin(L⋆−L, 0.03),L≤L⋆,max(0,Lref−L),L>L⋆,−10,invalid,r_t+1\;=\; cases (0,\,L_ref-L )\;+\;B\;+\;w\, (L -L,\;0.03 ),&L≤ L ,\\[3.0pt] (0,\,L_ref-L ),&L>L ,\\[3.0pt] -10,&invalid, cases (64) with B=5B=5 and w=100w=100. The relative reward is same as Equation 62: a valid attempt pays ρ^t+1δt+1 ρ_t+1\, _t+1, and an invalid one is charged ρ^t+1δmax ρ_t+1 _ . The reward rate is fitted on max(0,Lref−L) (0,L_ref-L) only, with the bonuses removed, so that a threshold bonus cannot inflate the price of a second. Vanilla. The baseline optimizes the benchmark’s own objective, the time to reach the target, and also similarly aims to reach further below the 3.283.28 target loss, with an exchange rate of w=20,000w=20,000 seconds per unit of loss. rt+1vanilla=−δt+1+wmin(L⋆−L, 0.03),L≤L⋆,−δmax−wmin(L−L⋆, 0.03),valid,L>L⋆,− 2δmax,invalid.r^\,vanilla_t+1\;=\; cases-\, _t+1\;+\;w\, (L -L,\;0.03 ),&L≤ L ,\\[3.0pt] -\, _ \;-\;w\, (L-L ,\;0.03 ),&valid,\;L>L ,\\[3.0pt] -\,2\, _ ,&invalid. cases (65) The three cases do not overlap, so every crossing attempt still outranks every non-crossing one and every valid attempt outranks every invalid one. The second case exists to encourage scripts that are near the target loss to cross the target. Buffer. The buffer is same as Section C.3.1 with the key of Equation 63 replaced by each arm’s own objective, r−ρ^δr- ρ\,δ for RPG and rvanillar^\,vanilla for the baseline. During selection from the buffer to form the current state sts_t, half of samples are reserved for the best samples in the buffer ranked by the key (for instance, since we have a batch size of 64, 32 samples would be selected this way). For the other half, it is ranked first by the number of changed lines and then by the key. The second half serves to encourage for exploration on making more significant code changes. C.4.4 Prompts and an Example Action See Figures 7, 8 and 9. Example Prompt for NanoGPT ⬇ You are competing in the NanoGPT speedrun (modded-nanogpt, track_1_short): train a neural network to at most 3.28 mean cross-entropy loss on the FineWeb validation set, using 8x NVIDIA A100-SXM4-80GB GPUs, in the least possible total wall-clock time. The score is the TOTAL WALL-CLOCK SECONDS OF THE WHOLE SCRIPT PROCESS, measured by our harness from process launch to process exit. You are scored on TWO things at once and must improve BOTH: the total wall-clock seconds, AND how far the final validation loss goes below 3.28 (worth roughly 20 seconds of wall-clock per 0.001 of loss, credited down to 3.25). Reaching 3.28 is required but not sufficient -- do not stop optimising the model the moment the script crosses it. A change that trades one against the other is NOT progress and is a wasted attempt: not a speedup paid for with loss margin, and not a loss improvement paid for with time. Propose the change that moves both in the right direction together. Below is an existing solution script and the result it achieved. Study it and propose ONE change that makes the whole process finish sooner AND drives the final validation loss further below 3.28. Both must improve; a proposal that helps one at the other’s expense does not count. You do NOT write code. You write a CHANGE SPECIFICATION, which a separate engineer will implement exactly as you describe it. The engineer implements ONLY what the CHANGE section says: if the CHANGE section is missing, vague, or names only a flag or a constant, that is what gets built. Your CHANGE section is your entire contribution. Hardware facts you must design for: A100 is Ampere (sm_80). It has NO FP8 tensor cores, so any FP8 code path (float8 matmuls, FP8 lm-head, FP8 scaling) is unavailable and must be avoided or replaced; bf16 and TF32 tensor cores are the fast paths. HBM2e at roughly 2 TB/s, NVLink between GPUs. FlashAttention-3 is Hopper-only and unavailable; FlashAttention-2, PyTorch SDPA, FlexAttention, Triton and torch.compile are available. Constraints the script must continue to satisfy: - Launched with: torchrun --nproc_per_node=8 train_gpt.py. Single self-contained file; it must not read any sibling source files. - Training data: data/fineweb10B/fineweb_train_*.bin relative to the working directory (equivalently os.environ.get("DATA_PATH", ".") + "/data/fineweb10B/"); validation: data/fineweb10B/fineweb_val_*.bin. Do not modify the train or validation token streams. Validation loss must be the mean cross-entropy over exactly 10485760 tokens of the validation set. - Each .bin shard has a 256-int32 header (magic 20240520, version, token count) followed by uint16 GPT-2 tokens; shard files are 100M tokens each. - Validation cadence is enforced: a validation pass exactly every 125 training steps (when step % 125 == 0) and once at the final step, each printing the standard val line; the validation loss computation and its printing must stay unchanged. Violating this scores the run as a failure. - Print these exact log line formats to stdout: after each training step, "step:<step>/<train_steps> train_time:<ms>ms step_avg:<ms>ms" (train_time as an integer number of milliseconds, step_avg with 2 decimals); at every validation, "step:<step>/<train_steps> val_loss:<loss> train_time:<ms>ms step_avg:<ms>ms" (val_loss with 4 decimals), for example: step:1390/1390 val_loss:3.2775 train_time:79532ms step_avg:57.22ms. The printed train_time is informational only; the score is the harness-measured wall-clock seconds of the whole process. - torch._inductor.config.coordinate_descent_tuning is banned. - No network access, no downloads. - The whole process (startup + compilation + training + every validation + exit) is killed at a hard wall cap of 1500 seconds and scored as a failure; compilation alone can take several minutes, so leave margin. Existing solution and its result: previous_plan_code Result of that solution: previous_plan_error What counts as a change. The change must be a MECHANISM, not a setting. Not acceptable, and treated as a wasted attempt: changing torch.compile arguments (mode, dynamic, fullgraph, backend), DDP arguments, or any hyperparameter on its own (batch size, sequence length, iteration count, learning rates, warmup/cooldown lengths, momentum, Newton-Schulz iteration count, gradient-accumulation factor); deleting or moving a line; renaming; changing where the printed timer starts. Acceptable: something that changes what the program computes or how data moves, for example a Triton fused kernel (RMSNorm, rotary, cross-entropy, Newton-Schulz orthogonalization, lm-head); a different attention backend or block-mask construction for sm_80; CUDA graphs or regional/selective compilation that removes recompiles; bf16/TF32 precision changes replacing an FP8 path; DDP bucket restructuring, sharded reduce-scatter/all-gather or communication overlap; async pinned H2D loading or loader prefetch; architecture changes (attention structure, value embeddings, U-net skips, logit softcap, sliding/varlen document-aware attention, MLP variants, embedding sparsification, parameter tying); optimizer mechanics (Muon/NorMuon variants, per-parameter-group treatment, cautious weight decay). Write your reply as exactly these three sections, in this order, and nothing else. Do not think out loud before the first section, do not restate the task or the script, and do not write out the whole file. Start immediately with the word DIAGNOSIS. DIAGNOSIS At most 120 words. What specifically is costing wall-clock time in the script above, and what evidence in the script or its result points at it. Be concrete about where. CHANGE At most 400 words. A numbered list of the exact edits: for each, the function, class, variable or line region to touch, what it should become, and exact values for any constant. New code must be described precisely enough that there is only one reasonable way to write it; short illustrative fragments are fine. Every edit must serve the ONE mechanism named in the first line of this section. RISK At most 60 words. What could break, and what the implementer should check to avoid it. Your reply is cut off at a hard token limit. A reply that has not reached the CHANGE section by then is worthless, so write the CHANGE section before anything long. Figure 7: Example prompt for NanoGPT. Judge Prompt for NanoGPT ⬇ You are the maintainer of the modded-nanogpt speedrun leaderboard, reviewing one submitted training script (train_gpt.py, launched with torchrun on 8x NVIDIA A100-SXM4-80GB GPUs). Your ONLY job is to decide whether it violates this rule: "Do not modify the train or validation data pipelines. You may change the batch size, sequence length, attention structure, model architecture, optimizer, precision, schedules, kernels, compilation, logging, timing -- anything -- just do not change the underlying streams of tokens, do not change what the printed validation loss MEANS, and do not change WHEN validation runs." The validation loss must remain: the mean cross-entropy of next-token prediction over EXACTLY ‘val_tokens‘ (10,485,760) tokens of the validation files (data/fineweb10B/fineweb_val_*.bin), read in their natural order, correctly averaged over everything that was summed (all validation batches AND all data-parallel ranks). The training stream must remain the tokens of data/fineweb10B/fineweb_train_*.bin, unmodified (order, content, and shard reading semantics; changing how MANY tokens are consumed per step via batch geometry is allowed). Validation must run exactly every 125 training steps (val_loss_every == 125, i.e. at step % 125 == 0) and once at the final step, printing the standard val line each time. VIOLATIONS include (non-exhaustive -- judge the semantics, not the syntax): * evaluating fewer validation tokens than val_tokens while still normalizing by the full count (e.g. the val loop no longer iterates over all val_steps, or val_steps is inflated relative to the batches actually evaluated); * removing or corrupting the cross-rank reduction of the validation loss (each rank sees 1/world_size of the tokens; the printed number must be the average over ranks); * dividing the validation loss by extra factors (world_size, an extra val_steps, a constant); * computing the "validation" loss on training data, synthetic tokens, a subset chosen by loss, or anything other than the validation files; * altering, filtering, reordering, or synthesizing the token streams; * printing a val_loss value that is not the number computed by the validation procedure above; * changing the validation cadence: val_loss_every different from 125, validating on a different schedule, skipping validations, validating only at the end, or adding / removing / moving the validation call (including gating it behind any condition other than the standard step % 125 == 0 or last-step check). NOT violations (do not flag these): * changed batch size, sequence length, gradient accumulation, number of training steps, early stopping of TRAINING when a target is reached; * architecture / attention / optimizer / LR schedule / precision / torch.compile changes, custom kernels; * harness bookkeeping: saving checkpoints to verify_ckpt/, build_model()/ eval_forward() helpers; * changes to the timing bookkeeping (where the clock starts / stops, what train_time prints): the printed time is informational only, the score is the harness-measured wall-clock of the whole process -- but the val line must remain parseable in its standard format; * refactors of the data loader that preserve which tokens are delivered in which order. Read the script carefully, especially the validation section and every line between computing per-batch losses and printing ‘val_loss:‘. Verify the arithmetic: what is summed, over how many batches, divided by what, reduced across ranks how. Then check the condition under which validation runs. If the script is so mangled it cannot run, judge only the pipeline rule (a crash is handled elsewhere). Reply with ONLY a JSON object, no other text: "verdict": "intact" | "modified", "reason": "<one sentence naming the exact line/mechanism if modified, or ’pipeline preserved’ if intact>" The submitted script: ‘python code ‘ Figure 8: Judge prompt for NanoGPT. Example Action for NanoGPT ⬇ DIAGNOSIS The script incurs massive startup and per-step overhead due to repeated ‘torch.compile‘ recompilation (only done once at start, but the model structure is complex), inefficient data loading via standard file I/O without async overlap, and suboptimal memory management where pinned tensors are allocated fresh per shard. The A100 lacks FP8, forcing full bf16/TF32 paths which are slower than Hopper’s FlashAttention-3 would be. The Muon optimizer adds significant compute cost per step for orthogonalization on gradients, slowing convergence and increasing wall-clock time per validation step. The current loader does not prefetch aggressively enough relative to the GPU bandwidth limits of 8x A100s, causing stalls during the initial training steps before warmup kicks in. CHANGE 1. Replace the sequential ‘_load_data_shard‘ call with an asynchronous pre-fetcher that loads the next two shards while the current one is being processed by the CPU/GPU pipeline, using ‘threading.Thread‘ or ‘concurrent.futures‘ to hide I/O latency behind the first micro-batch computation. Specifically, modify ‘DistributedDataLoader.advance()‘ to submit the next shard load immediately upon receiving the current batch, ensuring the CUDA stream remains busy. 2. Implement a custom Triton kernel for the Newton-Schulz orthogonalization step within ‘zeropower_via_newtrschulz5‘. This replaces the Python loop with a fused kernel that operates directly on pinned memory buffers passed from the optimizer, reducing Python interpreter overhead and enabling better occupancy on the A100’s tensor cores compared to the current CPU-bound gradient processing inside the optimizer step. 3. Introduce selective compilation: compile only the forward pass of the GPT model (excluding the backward pass) with ‘mode=’reduce-overhead’‘ and cache it globally, while keeping the backward pass uncompiled or compiled separately with ‘mode=’default’‘ to avoid recompiling the heavy gradient accumulation logic repeatedly if any dynamic shapes occur (though static here). Ensure the ‘compile_forward‘ function wraps only ‘model.forward‘ and is called exactly once before DDP wrapping. 4. Modify the training loop to perform a "warmup" phase where the sliding window starts at 1 block but the learning rate scheduler applies a linear ramp-up for the first 50 steps instead of immediate decay, allowing the network to stabilize without the aggressive orthogonalization pressure initially, effectively smoothing the loss landscape faster. 5. Add explicit ‘torch.cuda.synchronize()‘ calls only around the validation start and end points, removing them from the inner training loop to reduce synchronization overhead, relying on the harness’s global timing rather than per-step syncs that don’t contribute to the score. RISK Incorrectly implementing the async loader may cause race conditions where tokens are loaded out of order or double-loaded, breaking the validation token count requirement. The Triton kernel must strictly match the float precision (bf16) and output shape of the existing Python loop; mismatched dtypes will corrupt gradients. Ensure the warmup schedule doesn’t delay validation too much, violating the ‘step % 125 == 0‘ cadence. Verify that the new orthogonalization kernel handles edge cases like zero gradients gracefully to prevent NaN propagation. Figure 9: Example action for NanoGPT C.4.5 Hyperparameters See Table 3. Table 3: NanoGPT configuration. The upper block is shared by RPG and vanilla; the lower block is the complete set of settings on which they differ. Group Setting Value Model actor and reference Qwen3.5-4B critic Qwen3.5-4B, token-classification head rollout engine vLLM, bfloat16 PPO advantage estimator GAE, γ=1.0γ=1.0, λ=1.0λ=1.0 actor / critic learning rate 1×10−61× 10^-6 / 1×10−51× 10^-5 train batch == mini-batch 6464 clip ratio, value clip 0.20.2, 0.50.5 gradient clip 1.01.0 entropy coefficient 0.0010.001 KL coefficient to the reference 0.0010.001, fixed Rollout temperature, top-p, top-k 0.70.7, 0.80.8, 2020 presence penalty 1.51.5 samples per prompt 11 prompt / response length 40,96040,960 / 2,5602,560 tokens Environment wall-clock timeout δmax _ 15001500 s time floor 0.10.1 s target loss L⋆L , credited floor 3.283.28, 3.253.25 grader 8×8× A100 80GB per attempt validation cadence every 125125 steps Implementer model, temperature GPT 5.6 Terra, 00 response budget, attempts 16,38416,384 tokens, ≤3≤ 3 judge same model, 33 votes at temperature 00 Buffer selection width top 6464 large-edit quota 0.50.5 of slots, ≥20≥ 20 changed lines Compute trainer 4×4× A100 80GB graders 55 workers, 8×8× A100 80GB each repetitions 22 repetitions, each ran for 1010 steps RPG only reward reference LrefL_ref 1010 crossing bonus B, slope w 55, 100100 invalid reward −10-10 rate estimator online NIW, valid rows only forgetting factor, batches M 0.30.3, 1,0001,000 Vanilla only objective time to reach L⋆L depth exchange rate wtw_t 20,00020,000 s per unit loss invalid reward −2δmax-2\, _ Appendix D Experiments D.1 Bandit Experiments, ||=4|X|=4 D.1.1 Regret Curves Figure 2 reports only the endpoint of each run at T=30,000T=30,000. Figure 10 give the regret-per-step curves behind those endpoints. Figure 10: Regret against step for ||=4|X|=4 contexts. D.1.2 Hyperparameter Selection Every method that carries a hyperparameter is tuned via grid search. The tuning seeds are disjoint from the 100100 evaluation seeds, and selection is by median final regret over the tuning seeds. The hyperparameter search for every baseline are shown in Figures 11, 12 and 13. Figure 11: Tuning grid for C-UCB. The horizontal axis is c2c_2 and the vertical axis is c1c_1. Figure 12: Tuning grid for the NIW-NPG, over the learning rate and accumulation window. Figure 13: Tuning grid for the SPG-NIW, over the same learning rate and accumulation window grid as Figure 12. D.2 Bandit Experiments, ||>4|X|>4 D.2.1 Regret at the Horizon Figure 14 reports the endpoint (T=30,000T=30,000) of each run at ||∈8,16,32,64|X|∈\8,16,32,64\. (a) ||=8|X|=8. (b) ||=16|X|=16. (c) ||=32|X|=32. (d) ||=64|X|=64. Figure 14: Regret at the horizon against the number of arms K, one panel per environment family, for the approximate C-UCB setting. D.2.2 Regret Curves Figures 14 and 2 only contains regret at the endpoint T=30,000T=30,000. Figure 15 gives the regret per step at each context count. (a) ||=8|X|=8. (b) ||=16|X|=16. (c) ||=32|X|=32. (d) ||=64|X|=64. Figure 15: Regret against step for the large context setting, with environment families down the rows and arm counts across the columns. D.2.3 Hyperparameter Selection The tuning grids are produced by the same procedure as Section D.1.2, one grid per context count (Figures 16, 17 and 18). (a) ||=8|X|=8. (b) ||=16|X|=16. (c) ||=32|X|=32. (d) ||=64|X|=64. Figure 16: Tuning grid for approximate C-UCB. The horizontal axis is c2c_2 and vertical axis is c1c_1. (a) ||=8|X|=8. (b) ||=16|X|=16. (c) ||=32|X|=32. (d) ||=64|X|=64. Figure 17: Tuning grid for the NIW-NPG, over the learning rate and accumulation window. (a) ||=8|X|=8. (b) ||=16|X|=16. (c) ||=32|X|=32. (d) ||=64|X|=64. Figure 18: Tuning grid for the SPG-NIW, over the same learning rate and accumulation window grid as Figure 17. D.3 MLE Agent Competition identifiers. Table 1 abbreviates the MLE-Bench Lite competition ids to fit; in full they are: aerial-cactus-identification, aptos2019-blindness-detection, denoising-dirty-documents, dog-breed-identification, dogs-vs-cats-redux-kernels-edition, histopathologic-cancer-detection, detecting-insults-in-social-commentary, jigsaw-toxic-comment-classification-challenge, leaf-classification, mlsp-2013-birds, nomad2018-predict-transparent-conductors, new-york-city-taxi-fare-prediction, random-acts-of-pizza, plant-pathology-2020-fgvc7, ranzcr-clip-catheter-line-classification, siim-isic-melanoma-classification, spooky-author-identification, tabular-playground-series-dec-2021, tabular-playground-series-may-2022, text-normalization-challenge-english-language, text-normalization-challenge-russian-language, and the-icml-2013-whale-challenge-right-whale-redux. Figure 19: Per-task learning curves on MLE-Bench Lite, the evaluation method is same as Table 1. Table 1 reports each task at the final time budget, the figure shows the entire learning process. Figure 20: Mean cumulative time (x-axis) and mean cumulative reward (y-axis). RPG optimizes for long-term reward per unit of time, so its curve goes higher as mean cumulative time increases. Figure 21: Budget curve and an example self-improvement step on NanoGPT, with root script included for reference. (Left) At the final budget, RPG improves over vanilla by 85.7%85.7\%. Time 0 reflects the root parent script’s validation loss. We do not count the root script’s loss in the best crossings to reflect the dynamics of the policies. (Right) Example of a self-improvement edit by RPG, which rewrites Newton–Schulz in the Muon optimizer as a triton kernel with a compiled iteration and adds a learning-rate warmup. The new edit crosses the loss target using shorter time.