Paper deep dive
MLREF: Efficient Module Reuse for Reward Design in Reinforcement Learning via Large Language Models
Chenglin Liu, Xun Wang, Ruishuo Chen, Zhuoran Li, Longbo Huang
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 8/20/2026, 5:12:53 AM
Summary
The paper introduces MLREF (Module Level Reward Evolution Framework), a method for automated reward design in reinforcement learning using Large Language Models. Unlike existing methods that treat reward functions as monolithic programs, MLREF maintains a persistent 'module pool' of reusable reward components. The framework evolves this pool across iterations using reflection-based refinement, hybrid credit assignment, and a merge strategy with rollback. Experiments on 17 locomotion and manipulation tasks demonstrate that MLREF outperforms strong baselines like EUREKA and RF-Agent by 25.2% in locomotion and 6.6% in manipulation, while providing more stable optimization dynamics.
Entities (9)
Relation Signals (8)
MLREF → uses → Module Pool
confidence 95% · At the core of MLREF is a module pool, a persistent repository of reusable reward components.
MLREF → employs → Rollback Strategy
confidence 93% · MLREF integrates three mechanisms: reflection-based refinement, hybrid credit assignment, and a merge strategy with rollback
MLREF → employs → Hybrid Credit Assignment
confidence 93% · MLREF integrates three mechanisms: reflection-based refinement, hybrid credit assignment, and a merge strategy with rollback
MLREF → evaluatedon → Bi-DexHands
confidence 90% · We validate MLREF on 17 tasks from Isaac Gym and Bi-DexHands.
MLREF → evaluatedon → Isaac Gym
confidence 90% · We validate MLREF on 17 tasks from Isaac Gym and Bi-DexHands.
MLREF → outperforms → RF-Agent
confidence 90% · Experiments on 17 tasks show that MLREF outperforms strong baselines by 25.2% in locomotion
MLREF → outperforms → Eureka
confidence 90% · Experiments on 17 tasks show that MLREF outperforms strong baselines by 25.2% in locomotion
Large Language Models → enables → Automated Reward Generation
confidence 85% · While large language models (LLMs) have enabled automated reward generation
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Reward function design remains a bottleneck in reinforcement learning. While large language models (LLMs) have enabled automated reward generation, existing methods generate and revise reward functions as monolithic programs, making it difficult to reliably preserve and reuse effective components discovered in earlier iterations, leading to unstable performance across iterations. To address this, we propose Module Level Reward Evolution Framework (MLREF). At the core of MLREF is a module pool, a persistent repository of reusable reward components. MLREF treats the module pool as the primary optimization object: the pool evolves across iterations by accumulating successful modules, refining underperforming ones, and reusing proven components; while reward functions are constructed as linear combinations of modules drawn from this pool. To drive this evolution, MLREF integrates three mechanisms: reflection-based refinement, hybrid credit assignment, and a merge strategy with rollback, which together improve the effectiveness and robustness of reward optimization. Experiments on 17 tasks show that MLREF outperforms strong baselines by 25.2% in locomotion and 6.6% in manipulation, with more stable optimization dynamics.
Tags
Links
- Source: https://arxiv.org/abs/2608.18827v1
- Canonical: https://arxiv.org/abs/2608.18827v1
Trouble viewing inline? Open PDF directly →
Full Text
63,959 characters extracted from source content.
Expand or collapse full text
MLREF: Efficient Module Reuse for Reward Design in Reinforcement Learning via Large Language Models Chenglin Liu Xun Wang Ruishuo Chen Zhuoran Li Longbo Huang [0.3em]Institute for Interdisciplinary Information Sciences Tsinghua University Correspondence: longbohuang@tsinghua.edu.cn arXiv preprint, August 19, 2026 Abstract Reward function design remains a bottleneck in reinforcement learning. While large language models (LLMs) have enabled automated reward generation, existing methods generate and revise reward functions as monolithic programs, making it difficult to reliably preserve and reuse effective components discovered in earlier iterations, leading to unstable performance across iterations. To address this, we propose Module Level Reward Evolution Framework (MLREF). At the core of MLREF is a module pool, a persistent repository of reusable reward components. MLREF treats the module pool as the primary optimization object: the pool evolves across iterations by accumulating successful modules, refining underperforming ones, and reusing proven components; while reward functions are constructed as linear combinations of modules drawn from this pool. To drive this evolution, MLREF integrates three mechanisms: reflection-based refinement, hybrid credit assignment, and a merge strategy with rollback, which together improve the effectiveness and robustness of reward optimization. Experiments on 17 tasks show that MLREF outperforms strong baselines by 25.2% in locomotion and 6.6% in manipulation, with more stable optimization dynamics. 1. Introduction Reinforcement learning (RL) has achieved remarkable success in various domains 10; 38; 6; 25. However, its performance critically depends on the reward function, and designing an effective one remains a fundamental challenge. Sparse rewards are easy to define but often provide insufficient learning signals for optimization 26; 12, whereas dense rewards often require substantial domain expertise and manual engineering efforts 29; 7. Recent advances in large language models (LLMs) have provided new opportunities for automating reward design, leveraging their strong capabilities in reasoning 31; 37, instruction following 23, and code generation 2. However, directly generated rewards often suffer from hallucination, syntactic errors, and semantic misalignment with task objectives, limiting their reliability in practical RL scenarios 20; 27. These challenges motivate the development of iterative refinement frameworks, where LLMs progressively improve reward functions based on feedback from RL training. To improve the quality of LLM-generated rewards, numerous LLM-based harnesses and workflows have been developed to iteratively refine reward functions through interaction with RL training feedback. However, existing approaches predominantly optimize rewards at the function level by generating or modifying monolithic reward programs, e.g., EUREKA 20, R∗R^* 18 and RF-Agent 9. Without module-level tracking, credit assignment, and semantic recombination, these methods struggle to reuse effective reward components, resulting in redundant and inefficient search. Although promising, module-level reward optimization introduces several unique challenges beyond function-level optimization. In contrast to optimizing a complete reward program, module-level refinement requires identifying the semantic roles of individual reward components, determining their contribution to policy improvement, and effectively coordinating interactions among different modules. Moreover, maintaining a reward pool requires a tradeoff between preserving previously discovered effective modules and exploring new reward components, as naive accumulation may introduce redundancy or conflicting incentives. Therefore, designing a module-level reward optimization framework is a more challenging meta-level optimization problem: beyond leveraging holistic RL feedback for reward evaluation, it must attribute performance changes to individual modules and efficiently discover, refine, and reuse effective reward components. To address these limitations, we propose Module-Level Reward Evolution Framework (MLREF). At the core of MLREF is a module pool, a persistent repository consisting of multiple reusable reward components, each targeting a specific aspect of the task. This pool evolves across iterations by accumulating successful modules, refining underperforming ones, and reusing proven components for future reward construction. At each iteration, reward functions are assembled as linear combinations of modules drawn from the pool. The framework integrates three mechanisms: (i) reflection-based refinement that leverages task analysis and historical feedback to guide module design; (i) hybrid weight optimization that empirically assigns credit to individual modules; and (i) a merge strategy with rollback that consolidates successful modules across parallel samples while recovering from failures. Together, these mechanisms enable stable, empirically grounded pool evolution: each module’s design is informed by task analysis and historical feedback, its contribution is credited empirically rather than assumed, and successful modules are consolidated while failures are rolled back, so the reward improves steadily rather than oscillating across iterations. Table 1 summarizes the key differences between MLREF and prior LLM-based reward design methods. Property EUREKA RF-Agent MLREF (Ours) Optimization Unit Function Function Module Pool Module Identity – Anonymous Named + Spec. Cross-iter. Persistence – Partial Persistent Pool Credit Assignment Holistic Holistic Per-Module Hybrid Failure Recovery – Implicit Explicit Rollback Semantic Guidance – – Task + Feedback Refl. Table 1: Comparison of key design properties across LLM-based reward design methods. MLREF is the only framework that treats modules as first-class, persistently managed objects with per-module credit assignment, semantic guidance, and a principled failure recovery mechanism. We validate MLREF on 17 tasks from Isaac Gym and Bi-DexHands. Experiments show that MLREF outperforms state-of-the-art LLM-based reward design methods by 25.2% in locomotion and 6.6% in manipulation, and ablation studies verify the contributions of individual components. Evolution analysis further reveals that MLREF maintains stable optimization trajectories while baselines suffer from severe performance oscillation, confirming the effectiveness of the rollback mechanism in achieving iterative stability. In summary, our contributions are: • We identify the inability of function-level optimization to efficiently reuse modules as a key driver of performance oscillation in LLM-based reward evolution. • We introduce the module pool abstraction, which shifts the optimization object from individual reward functions to a persistent repository of reusable modules, and develop three complementary mechanisms that together enable stable, directed pool evolution. • Across 17 locomotion and manipulation tasks, MLREF sets a new state of the art and, more importantly, exhibits substantially more stable optimization dynamics than prior methods, as confirmed by both ablation and evolution analysis. 2. Related Work Before the advent of LLMs, reward design primarily relied on inverse RL 1 and preference-based RL 4; both require substantial human effort and generalize poorly across tasks. Early LLM-based work used LLMs as black-box generators queried per timestep 14; 5, which is computationally prohibitive, or generated code-form rewards without iterative refinement 36. Later methods introduced iterative refinement guided by human feedback on visualized behaviors 33; 11; 13, but still depend on manual intervention, limiting scalability. EUREKA 20 established the first fully automated pipeline for iterative reward optimization. Building on this, R* 18 introduced AST-based crossover with preference-based parameter tuning; CARD 27 applied heuristic pre-screening before costly RL training; RF-Agent 9 adapted MCTS for reward space search; and FORGE 8 incorporated chain-of-thought reasoning and systematic crossover. LLM-based reward design has also been explored in game-playing 32; 16; 35, multi-agent RL 17; 19; 30, and offline imitation learning 28. None of these methods, however, maintain an explicit module pool for systematic reuse or account for LLM-induced variance in evaluation. In contrast, MLREF is the first framework to maintain a persistent module pool for systematic module-level accumulation, refinement, and reuse throughout the reward evolution process. 3. Preliminaries 3.1. Markov Decision Process and Reinforcement Learning We consider sequential decision-making problems formalized as Markov Decision Processes (MDPs) 24, defined by (,,P,R,γ,ρ0)(S,A,P,R,γ, _0) with state space S, action space A, transition kernel P, reward function R:×→ℝR:S×A , discount factor γ, and initial state distribution ρ0 _0. A reinforcement learning algorithm seeks a policy π:→Δ()π:S→ (A) that maximizes the expected discounted return J(π,R)=π,P,ρ0[∑tγtR(st,at)]J(π;R)=E_π,P, _0\! [ _tγ^tR(s_t,a_t) ]. Because R is the sole task-specific signal in this objective, the choice of reward function entirely determines the behavior that is learned. 3.2. Reward Design Problem Based on the MDP defined above, a reward design problem (RDP) 15 is characterized by the tuple ⟨ℰ,ℛ,,F⟩ ,R, A,F , where ℰ=(,,P)E=(S,A,P) denotes the environment, ℛR is the space of candidate reward functions, A is the policy optimization algorithm, and F:Π→ℝF: is the fitness function that evaluates the task performance of a policy. For each reward function R∈ℛR , let πR=(ℰ,R) _R= A(E,R) denote the policy obtained by optimizing R in ℰE. The goal of an RDP is to identify a reward function whose resulting policy maximizes the fitness: R∗=argmaxR∈ℛF(πR,ℰ)R^*= _R F( _R;E). 4. Method 4.1. Overview Inspired by the prior work 20, MLREF models a reward function as a linear combination of reward modules, where each module targets a specific aspect of the task and, when combined, forms a robust reward function that balances multiple desiderata. Formally, a reward function R is constructed from a set of K modules m1,m2,…,mK\m_1,m_2,…,m_K\ with corresponding weights w1,w2,…,wK\w_1,w_2,…,w_K\: R(s,a)=∑k=1Kwk⋅mk(s,a).R(s,a)= _k=1^Kw_k· m_k(s,a). (1) To fully leverage the potential of these modules, MLREF introduces a module pool, a persistent repository where modules are accumulated, refined, and recombined across iterations. Figure 1 illustrates the full pipeline and Algorithm 1 presents the corresponding procedure; the pipeline consists of two phases: initialization and iterative optimization. In the initialization phase, the LLM performs initial reflection on the task description and environment code (steps 1–2), then generates S diverse module pool variants (step 3). In the iterative optimization phase, each variant constructs a reward function via weighted linear combination (step 4), which is evaluated through RL training (step 5). A pool merge step consolidates modules from successful variants and rolls back upon failure based on the training feedback (steps 6–7). The LLM then conducts feedback reflection to analyze training outcomes and generates refined pool variants for the next iteration (steps 8–9). Complete LLM prompt templates are provided in the Technical Supplement. Figure 1: Overview of the MLREF pipeline. In the initialization phase (steps 1–3), the LLM performs task and environment reflection and generates diverse module pool variants. In the iterative optimization phase (steps 4–9), each variant assembles a reward function via weighted linear combination for RL evaluation; a merge strategy with rollback consolidates successful modules; and feedback reflection guides the next round of pool improvement. The key departure from prior work lies in what is being optimized. In MLREF, a reward function is only a temporary instance assembled from the pool for RL evaluation; the pool itself is the persistent object that accumulates evidence and improves over iterations. As a result, effective modules persist beyond the function that introduced them and are actively reused in subsequent iterations; training feedback is applied at the module level, enabling finer-grained and more targeted optimization than function-level rewriting. Algorithm 1 MLREF Iterative Optimization 0: Task description T, environment code E, LLM ℳM, RL trainer A, merge strategy Merge, max iterations I, parallel samples S 0: Best reward function R∗R^* global←∅P_global← , prev←∅P_prev← , best←∅S_best← , R∗←∅R^*← for iter=1iter=1 to I do // Stage 1: Module pool generation or improvement if iter=1iter=1 then ref←ℳ.Reflect(T,E)ref . Reflect(T,E) // Initial Reflection for i=1i=1 to S in parallel do i←ℳ.InitPool(T,E,ref)P_i . InitPool(T,E,ref) // Pool Initialization end for else ref←ℳ.Reflect(prev,best,prev,global)ref . Reflect(S_prev,S_best,P_prev,P_global) // Feedback Reflection for i=1i=1 to S in parallel do i←ℳ.ImprovePool(T,E,ref,global)P_i . ImprovePool(T,E,ref,P_global) // Pool Improvement end for end if // Stage 2: Reward construction and RL evaluation for i=1i=1 to S in parallel do Ri←ConstructReward(i)R_i← ConstructReward(P_i) // Weighted linear combination (Eq. 1) i←.Train(E,Ri)S_i . Train(E,R_i) end for // Stage 3: Pool merge and best selection global←Merge(i,i)P_global← Merge(\P_i\,\S_i\) // Accumulate or rollback (prev,prev)←argmaxii(S_prev,P_prev) _i\;S_i // Select the best pool if prev>bestS_prev>S_best then best←prevS_best _prev, R∗←corresponding reward of prevR^* reward of S_prev end if end for return R∗R^* 4.2. Module Pool Below we describe the three stages of pool evolution: initialization, improvement, and merge. 4.2.1 Pool Initialization In the first iteration, MLREF constructs a module pool from scratch. The initialization proceeds in two steps: specification generation and module implementation. A module’s specification is its metadata, consisting of the module name, input variable names and types, and a natural-language description of its intended function. In the specification generation step, the LLM generates a set of module specifications that collectively cover the diverse requirements of the task, guided by the task description, environment code, and initial reflection. In the implementation step, the LLM generates Python code for each module based on its specification. Separating specification from implementation improves generation quality by letting the LLM focus on one subtask per response. 4.2.2 Pool Improvement From the second iteration onward, MLREF refines the existing module pool. Improvement also follows a two-step structure: improvement plan design and improvement plan execution. In the plan design step, the LLM produces an improvement plan based on the training statistics and feedback reflection. We allow four types of operations on modules: 1. Add. Introduce a new module by providing its specification (analogous to initialization). 2. Delete. Remove a module that is redundant or empirically harmful. 3. Modify. Revise a module’s code while keeping its specification fixed (e.g., fixing incorrect structure or unsuitable output scale). 4. Rewrite. Replace both the specification and the code when the specification itself is deficient. For each operation, the LLM also provides a rationale. In the execution step, the LLM implements the concrete code changes relating to new modules, modified modules and rewritten modules according to the plan. This two-step decomposition again improves output quality by separating planning from coding. By operating at the per-module level, MLREF achieves finer-grained control over reward function optimization than prior function-level approaches. 4.2.3 Pool Merge At each iteration, MLREF produces S module pool variants in parallel and consolidates them via a module-wise merge. First, variants whose RL performance falls below the historical best by more than a margin t are discarded. All modules from surviving variants are then collected into a single set; when multiple variants contain modules with the same name, the best-performing version is retained. Since all variants derive from the same base pool, this merge consolidates diverse improvement plans without sacrificing quality. If no variant meets the threshold, MLREF rolls back to the previous pool. This rollback mechanism is central to MLREF’s iterative stability: it prevents a single poor iteration from destroying the entire optimization trajectory. The failed attempts are recorded and fed into the next round of feedback reflection, encouraging the LLM to explore different strategies. 4.3. Reflection Despite decomposing complex processes into multiple steps, LLM outputs can still suffer from hallucination (producing invalid code that crashes RL training) or insufficient reasoning (yielding suboptimal module designs). To improve output stability and quality, we introduce a reflection mechanism inspired by prior work 19; 28 using chain-of-thought prompting. As shown in Fig. 1 (steps 2 and 8), MLREF employs two types of reflection: initial reflection before the first iteration, and feedback reflection before each subsequent improvement. 4.3.1 Initial Reflection Before the first pool initialization, MLREF conducts initial reflection in two parts. Task reflection prompts the LLM to analyze the task description, identifying conditions for success, behaviors likely to help or hinder performance, and non-obvious strategies worth exploring. We explicitly encourage diversity and forbid concrete code generation at this stage. Environment reflection extracts relevant observation variables including their names, types, tensor shapes, and descriptions from the environment code, which provides a reference for downstream module design and a validity check against undefined variables. Task and environment reflections are separated to ensure depth in each, and the resulting analysis is reused during all subsequent pool improvement steps. 4.3.2 Feedback Reflection From the second iteration onward, MLREF performs feedback reflection on the previous round’s training outcome before planning pool improvements. If training completed successfully, the LLM receives the historical best pool (with its module composition and training results) alongside the previous round’s pool, improvement plan, and results. It analyzes reward trends per module, compares module design and composition between the two rounds, and attributes performance changes to specific pool differences. When the previous round achieves a new best, the LLM extracts successful design patterns to reinforce; when performance regresses, it diagnoses weaknesses in the improvement plan and proposes new directions. If training encountered an error, the LLM instead receives the error trace, the previous improvement plan, and the environment code. It locates the error source, analyzes its cause, and suggests concrete measures to avoid similar failures in future iterations. As with initial reflection, we encourage the LLM to explore diverse hypotheses and forbid it from generating concrete improvement plans or code during this step, preserving a clean separation between analysis and execution. 4.4. Weight Optimization Assigning appropriate weights to modules is critical: different modules vary in importance and output magnitude, so uniform weighting is suboptimal, while relying solely on the LLM to assign weights leads to inconsistency across iterations. MLREF addresses this with a hybrid weight optimization strategy that combines LLM-based semantic judgment with empirical evidence from RL training. Each module maintains two credit scores. LLM credit captures the LLM’s assessment of a module’s relevance: after pool initialization or improvement, the LLM is prompted to select modules and propose initial weights, which are normalized to form LLM credits (with unselected modules receiving zero). Correlation credit captures the module’s empirical contribution: we compute the Pearson correlation between the module’s reward sequence and the performance curve during training, incorporating a time-lag compensation to account for delayed effects of rewards on performance. Both the raw reward–performance correlation and the correlation of their temporal differences are combined into a single correlation score. To prevent credit values from oscillating across iterations, both LLM and correlation credits are smoothed via exponential moving average. Before combining the two credits, we normalize them to comparable ranges: linear normalization for the non-negative LLM credit, and softmax normalization for the signed correlation credit (amplifying the distinction between positively and negatively correlated modules). The normalized scores are then fused into a composite score via a weighted sum. For module selection, we apply an upper confidence bound (UCB) term over the composite score to balance exploitation of high-scoring modules with exploration of underused ones. The top-K modules by UCB score are selected, and their composite scores are used directly as the weights wkw_k in the linear combination (Eq. 1). The full mathematical formulation is provided in the Technical Supplement. 5. Experiments 5.1. Experimental Setup We evaluate MLREF on 17 representative tasks from Isaac Gym 21 and Bi-DexHands 3, covering locomotion and dexterous manipulation challenges (see Table 2 for the full list) to answer the following questions: (i) How does MLREF perform compared to state-of-the-art baselines? (i) How do individual components of MLREF contribute to its performance? (i) How stable is the reward evolution process under MLREF? We compare against two state-of-the-art baselines: EUREKA 20 generates reward function variants at each iteration by prompting the LLM to improve upon the previous best candidate, using the task description, environment code, and training feedback. We follow the original configuration: 5 iterations with 16 samples per iteration. RF-Agent 9 maintains a Monte Carlo tree over the reward space, selecting promising nodes for targeted optimization, evaluating them via RL, and back-propagating the results. We set the number of MCTS updates to 80, matching the total RL training budget of EUREKA. For MLREF, to mitigate the randomness introduced by the LLM, we run each task for 3 independent complete pipeline executions and return the reward function from the best execution. Each execution consists of I=9I=9 iterations with S=3S=3 parallel module pool variants per iteration, yielding a total RL training budget comparable to the baselines. For each algorithm on each task, RL training runs for 3,000 epochs within each iteration. After the final reward function is produced, we evaluate it by training a policy from scratch for 20,000 epochs using PPO with 5 independent random seeds, following the RL configuration of EUREKA 20. All LLM calls use the DeepSeek-V4-Flash model 34 in reasoning mode via the official API. Full hyperparameter configurations and LLM prompt templates are provided in the Technical Supplement. The environment setting follows the configuration of EUREKA 20. Task EUREKA RF-Agent MLREF (Ours) Locomotion Allegro Hand 25.48 ± 1.27 27.41 ± 1.90 28.89 ± 1.36 Ant 12.35 ± 1.29 12.34 ± 0.70 12.47 ± 0.30 Anymal -0.0057 ± 0.0021 -0.0064 ± 0.0010 -0.0073 ± 0.0019 Franka Cabinet 0.397 ± 0.202 0.701 ± 0.200 0.997 ± 0.006 Humanoid 7.95 ± 0.87 5.56 ± 1.63 7.40 ± 1.12 Quadcopter -0.042 ± 0.007 -0.277 ± 0.485 -0.040 ± 0.011 Shadow Hand 17.11 ± 2.17 16.53 ± 1.65 11.75 ± 9.69 Avg. Normalized Score 2.135 2.625 3.288 (+25.2%) Manipulation Block Stack 0.116 ± 0.050 0.139 ± 0.096 0.313 ± 0.244 Bottle Cap 0.994 ± 0.004 0.992 ± 0.010 0.987 ± 0.012 Catch Abreast 0.718 ± 0.066 0.657 ± 0.047 0.644 ± 0.059 Catch Underarm 0.910 ± 0.010 0.909 ± 0.016 0.871 ± 0.033 Door Close Outward 0.266 ± 0.129 0.415 ± 0.265 0.256 ± 0.084 Grasp and Place 0.710 ± 0.345 0.182 ± 0.130 0.475 ± 0.040 Kettle 0.441 ± 0.376 1.000 ± 0.000 0.901 ± 0.198 Lift Underarm 0.569 ± 0.465 0.512 ± 0.394 0.930 ± 0.022 Over 0.968 ± 0.007 0.985 ± 0.002 0.966 ± 0.011 Swing Cup 0.945 ± 0.093 0.721 ± 0.363 0.736 ± 0.294 Avg. Performance 0.664 0.651 0.708 (+6.6%) Table 2: Performance comparison across 17 tasks. Each row shows the mean and standard deviation of the raw performance metric. The score of locomotion tasks is normalized before averaging. The percentage improvement is relative to the best baseline. Bold indicates the best mean result. 5.2. Main Results To answer Q1, table 2 reports the performance of MLREF compared to EUREKA and RF-Agent across all 17 tasks. We report the mean and standard deviation of the raw performance metrics. For locomotion tasks, we further report the normalized score raw−sparsehuman−sparse raw-sparsehuman-sparse, where sparse and human baselines are taken from RF-Agent 9; higher values indicate better performance. Bold indicates the best mean result. Across the 17 tasks, MLREF achieves the best average performance in both categories. For locomotion, MLREF attains the highest mean on 4 of 7 tasks, with an average normalized score of 3.288, representing a 25.2% improvement over the best baseline. The gain is most striking on Franka Cabinet, where MLREF approaches the maximum possible score (0.997 vs. 0.701 of the nearest baseline). For manipulation, although MLREF obtains the best individual result on only 2 of 10 tasks, its average performance (0.708) is the highest among all methods, yielding a 6.6% improvement over the best baseline, with particularly strong results on Block Stack and Lift Underarm. These results demonstrate that systematic module-level management via the module pool leads to more effective reward functions than function-level optimization approaches. 5.3. Ablation Study Method Ant Block Stack Bottle Cap Avg. % of Full Full MLREF 12.47 ± 0.30 0.313 ± 0.244 0.987 ± 0.012 100.0% GPT-4o 10.77 ± 1.95 0.101 ± 0.042 0.041 ± 0.027 40.9% No Pool 11.51 ± 0.64 0.085 ± 0.040 0.988 ± 0.010 73.2% No Reflection 9.58 ± 0.55 0.057 ± 0.032 0.639 ± 0.401 53.3% No Weight Opt. 11.68 ± 1.14 0.044 ± 0.025 0.829 ± 0.338 63.9% Table 3: Ablation results on three representative tasks. Best results are bold. To answer Q2, we conduct ablation experiments on three representative tasks (Ant, Bottle Cap, Block Stack) by removing or replacing individual mechanisms from the full MLREF framework: 1. GPT-4o: replace the LLM backbone (DeepSeek-V4-Flash) with GPT-4o model 22. 2. No Pool: remove the module pool; at each iteration, the LLM directly generates a complete reward function, analogous to EUREKA. 3. No Reflection: remove both initial and feedback reflection; the LLM generates and improves pools directly from the task description and raw training feedback. 4. No Weight Opt.: remove the hybrid weight optimization strategy; the LLM selects modules and assigns weights without empirical credit assignment. Table 3 reports the results. All other experimental settings remain identical to the full MLREF configuration. The full MLREF framework achieves the best or near-best performance across all tasks. Replacing the LLM backbone with GPT-4o causes the most severe degradation (Avg. 40.9%), though MLREF still improves over iterations, indicating that pool evolution can accumulate effective modules even from a weak initialization. Disabling reflection incurs the most consistent degradation (Avg. 53.3%), highlighting its critical role in stabilizing optimization. The No Pool variant matches the full framework on Bottle Cap (0.988 vs. 0.987) due to serendipitous exploration, but without the pool, successful designs cannot be stably inherited. We also observed that only the full MLREF framework consistently generalized from validation to test: across all three tasks, the reward selected by its highest validation score remained the best at test time. In contrast, the ablated variants frequently produced rewards that scored well during validation but underperformed at test, suggesting a tendency to overfit when individual components are removed. 5.4. Iterative Dynamics Figure 2: Performance evolution across iterations (best sample per iteration). Left: all-task average. Middle: Shadow Hand. Right: Catch Abreast. Shaded areas denote standard deviation. To answer Q3, Figure 2 presents the performance evolution across optimization iterations on average (left panel) and for two representative tasks (middle and right panels). For each algorithm, we track the performance of the best sample at each iteration: for MLREF, this is the best among all parallel pool variants across 3 independent runs; for RF-Agent, the best node expanded at each MCTS step; for EUREKA, the best among the 16 samples of that iteration. Since the three methods operate with different iteration counts, we normalize the horizontal axis to [0,1][0,1] according to evolution progress. For the all-task average, we min-max normalize the per-task performance to ensure that each task contributes equally to the aggregate curve. On the all-task average, EUREKA benefits from large sample counts in early iterations but plateaus thereafter, while MLREF, starting from a lower initial point, exhibits consistent improvement through modular directed optimization; RF-Agent’s MCTS-based strategy yields intermediate behavior. On Shadow Hand and Catch Abreast, baselines suffer from severe performance oscillation. In particular, EUREKA degrades sharply after a strong first iteration on Catch Abreast. By contrast, MLREF maintains stable trajectories through its rollback mechanism while continuing to improve via iterative refinement, demonstrating strong optimization stability. Complete evolution curves for all 17 tasks are provided in the Technical Supplement. 5.5. Discussion Our results demonstrate the effectiveness of module-level management. Several limitations suggest directions for future work. First, on sparse-reward manipulation tasks such as Block Stack, MLREF’s absolute performance remains modest, suggesting stronger early exploration and progressive exploitation could improve the balance. Second, extending the evaluation of MLREF to a broader range of LLMs, including open-source models, would further validate its generality. Finally, while MLREF currently operates with a single LLM, multi-LLM collaboration to leverage diverse reasoning and knowledge represents a promising direction. 6. Conclusion We presented MLREF, a module pool-based framework that optimizes reusable reward modules through systematic pool-level operations. Unlike prior function-level approaches, MLREF’s module pool enables persistent accumulation, refinement, and reuse of reward components across iterations, supported by reflection, hybrid credit assignment, and a rollback-equipped merge strategy that together achieve iterative stability. Experiments on 17 tasks demonstrate that MLREF outperforms state-of-the-art LLM-based reward design methods, achieving 25.2% average improvement in locomotion and 6.6% in manipulation. Ablation studies confirm the contributions of each component, with reflection being particularly critical. Evolution analysis further shows that MLREF maintains stable optimization trajectories through its rollback mechanism. References Adams et al. (2022) S. Adams, T. Cody, and P. A. Beling A survey of inverse reinforcement learning. Artificial Intelligence Review 55 (6), p. 4307–4346. Cited by: §2. Chen et al. (2021) M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. D. O. Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374. Cited by: §1. Chen et al. (2022) Y. Chen, T. Wu, S. Wang, X. Feng, J. Jiang, Z. Lu, S. McAleer, H. Dong, S. Zhu, and Y. Yang Towards human-level bimanual dexterous manipulation with reinforcement learning. In Advances in Neural Information Processing Systems, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (Eds.), Vol. 35, p. 5150–5163. External Links: Link Cited by: §5.1. Christiano et al. (2017) P. F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei Deep reinforcement learning from human preferences. Advances in neural information processing systems 30. Cited by: §2. Du et al. (2023) Y. Du, O. Watkins, Z. Wang, C. Colas, T. Darrell, P. Abbeel, A. Gupta, and J. Andreas Guiding pretraining in reinforcement learning with large language models. In International Conference on Machine Learning, p. 8657–8677. Cited by: §2. Elguea-Aguinaco et al. (2023) Í. Elguea-Aguinaco, A. Serrano-Muñoz, D. Chrysostomou, I. Inziarte-Hidalgo, S. Bøgh, and N. Arana-Arexolaleiba A review on reinforcement learning for contact-rich robotic manipulation tasks. Robotics and Computer-Integrated Manufacturing 81, p. 102517. Cited by: §1. Eschmann (2021) J. Eschmann Reward function design in reinforcement learning. In Reinforcement Learning Algorithms: Analysis and Applications, p. 25–33. External Links: ISBN 978-3-030-41188-6, Document, Link Cited by: §1. Fan and Du (2025) H. Fan and J. Du Forging better rewards: a multi-agent LLM framework for automated reward evolution. External Links: Link Cited by: §2. Gao et al. (2026) N. Gao, X. Zhang, X. Jiang, M. You, M. Zhang, and Y. Deng Rf-agent: automated reward function design via language agent tree search. Advances in Neural Information Processing Systems 38, p. 172532–172577. Cited by: Table 4, §1, §2, §5.1, §5.2. Goldwaser and Thielscher (2020) A. Goldwaser and M. Thielscher Deep reinforcement learning for general game playing. In Proceedings of the AAAI conference on artificial intelligence, Vol. 34, p. 1701–1708. Cited by: §1. Guo et al. (2024) Q. Guo, X. Liu, J. Hui, Z. Liu, and P. Huang Utilizing large language models for robot skill reward shaping in reinforcement learning. In International Conference on Intelligent Robotics and Applications, p. 3–17. Cited by: §2. Hare (2019) J. Hare Dealing with sparse rewards in reinforcement learning. External Links: 1910.09281, Link Cited by: §1. Hazra et al. (2025) R. Hazra, A. Sygkounas, A. Persson, A. Loutfi, and P. Zuidberg Dos Martires REvolve: reward evolution with large language models using human feedback. In International Conference on Learning Representations, Vol. 2025, p. 101949–101990. Cited by: §2. Kwon et al. (2023) M. Kwon, S. M. Xie, K. Bullard, and D. Sadigh Reward design with language models. arXiv preprint arXiv:2303.00001. Cited by: §2. Lewis et al. (2010) R. L. Lewis, S. Singh, and A. G. Barto Where do rewards come from. In Proceedings of the international symposium on AI-inspired biology, p. 2601–2606. Cited by: §3.2. Li et al. (2024) H. Li, X. Yang, Z. Wang, X. Zhu, J. Zhou, Y. Qiao, X. Wang, H. Li, L. Lu, and J. Dai Auto mc-reward: automated dense reward design with large language models for minecraft. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 16426–16435. Cited by: §2. Li et al. (2025a) L. Li, L. Yuan, P. Liu, T. Jiang, and Y. Yu LLM-assisted semantically diverse teammate generation for efficient multi-agent coordination. In Forty-second International Conference on Machine Learning, Cited by: §2. Li et al. (2025b) P. Li, H. Jianye, H. Tang, Y. Yuan, J. Qiao, Z. Dong, and Y. Zheng R*: efficient reward design via reward structure evolution and parameter alignment optimization with large language models. In Forty-second International Conference on Machine Learning, Cited by: §1, §2. Li et al. (2025c) P. Li, H. Tang, Y. Yuan, and J. HAO ReMAC: large language model-driven reward design for multi-agent manipulation collaboration. In Workshop on Scaling Environments for Agents, External Links: Link Cited by: §2, §4.3. Ma et al. (2024) Y. J. Ma, W. Liang, G. Wang, D. Huang, O. Bastani, D. Jayaraman, Y. Zhu, J. Fan, et al. Eureka: human-level reward design via coding large language models. In International conference on learning Representations, Vol. 2024, p. 26516–26560. Cited by: Table 4, Table 4, Table 4, §1, §1, §2, §4.1, §5.1, §5.1, §5.1. Makoviychuk et al. (2021) V. Makoviychuk, L. Wawrzyniak, Y. Guo, M. Lu, K. Storey, M. Macklin, D. Hoeller, N. Rudin, A. Allshire, A. Handa, and G. State Isaac gym: high performance gpu-based physics simulation for robot learning. External Links: 2108.10470, Link Cited by: §5.1. OpenAI (2024) OpenAI Hello gpt-4o. Note: https://openai.com/index/hello-gpt-4o/ Cited by: item 1. Ouyang et al. (2022) L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems 35, p. 27730–27744. Cited by: §1. Puterman (1990) M. L. Puterman Markov decision processes. Handbooks in operations research and management science 2, p. 331–434. Cited by: §3.1. Radosavovic et al. (2024) I. Radosavovic, T. Xiao, B. Zhang, T. Darrell, J. Malik, and K. Sreenath Real-world humanoid locomotion with reinforcement learning. Science Robotics 9 (89), p. eadi9579. Cited by: §1. Stanton and Clune (2018) C. Stanton and J. Clune Deep curiosity search: intra-life exploration improves performance on challenging deep reinforcement learning problems. arXiv preprint arXiv:1806.00553. Cited by: §1. Sun et al. (2025a) S. Sun, R. Liu, J. Lyu, J. Yang, L. Zhang, and X. Li A large language model-driven reward design framework via dynamic feedback for reinforcement learning. Knowledge-Based Systems 326, p. 114065. Cited by: §1, §2. Sun et al. (2025b) S. Sun, J. Lyu, R. Liu, M. Yan, B. Liu, D. Ye, and X. Li PROF: an llm-based reward code preference optimization framework for offline imitation learning. arXiv preprint arXiv:2511.13765. Cited by: §2, §4.3. Sutton and Barto (1998) R. S. Sutton and A. G. Barto Reinforcement learning: an introduction. MIT Press, Cambridge, MA, USA. External Links: ISBN 0-262-19398-1, Link Cited by: §1. Wei et al. (2026) D. Wei, P. Yi, J. Lei, Y. Hong, H. Dong, and Y. Du An automated reinforcement learning reward design framework with large language model for cooperative platoon coordination. IEEE Transactions on Intelligent Transportation Systems. Cited by: §2. Wei et al. (2022) J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, Q. V. Le, D. Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems 35, p. 24824–24837. Cited by: §1. Wu et al. (2023) Y. Wu, Y. Fan, P. P. Liang, A. Azaria, Y. Li, and T. M. Mitchell Read and reap the rewards: learning to play atari with the help of instruction manuals. Advances in Neural Information Processing Systems 36, p. 1009–1023. Cited by: §2. Xie et al. (2024) T. Xie, S. Zhao, C. Wu, Y. Liu, Q. Luo, V. Zhong, Y. Yang, and T. Yu Text2reward: reward shaping with language models for reinforcement learning. In International Conference on Learning Representations, Vol. 2024, p. 35663–35699. Cited by: §2. Xu et al. (2026) A. Xu, B. Lin, B. Xue, B. Wang, B. Xu, B. Wu, B. Zhang, C. Lin, C. Dong, C. Ling, et al. Deepseek-v4: towards highly efficient million-token context intelligence. arXiv preprint arXiv:2606.19348. Cited by: §5.1. Yifan et al. (2025) H. Yifan, H. Bin-Bin, Y. Bowen, and Z. Hai-Tao LLM coach: reward shaping for reinforcement learning-based navigation agent. In 2025 Joint International Conference on Automation-Intelligence-Safety (ICAIS) & International Symposium on Autonomous Systems (ISAS), p. 1–6. Cited by: §2. Yu et al. (2023) W. Yu, N. Gileadi, C. Fu, S. Kirmani, K. Lee, M. G. Arenas, H. L. Chiang, T. Erez, L. Hasenclever, J. Humplik, et al. Language to rewards for robotic skill synthesis. arXiv preprint arXiv:2306.08647. Cited by: §2. Zhang et al. (2024) X. Zhang, C. Du, T. Pang, Q. Liu, W. Gao, and M. Lin Chain of preference optimization: improving chain-of-thought reasoning in llms. Advances in Neural Information Processing Systems 37, p. 333–356. Cited by: §1. Zhu and Zhang (2021) K. Zhu and T. Zhang Deep reinforcement learning based mobile robot navigation: a review. Tsinghua Science and Technology 26 (5), p. 674–691. Cited by: §1. Appendix Appendix A Weight Optimization Details This section provides the mathematical formulation of the hybrid weight optimization strategy discussed in the main paper. We describe the pipeline in three stages: credit assignment, credit normalization and fusion, and module selection. A.1. Credit Assignment MLREF maintains two credit scores per module: LLM credit and correlation credit. A.1.1 LLM Credit After pool initialization or improvement, the LLM is prompted to select modules and propose initial weights. The assigned weights are normalized to sum to 1, serving as the LLM credit ckLLMc_k^LLM for the current iteration; unselected modules receive zero credit. A.1.2 Correlation Credit During RL training in the previous iteration, let the performance curve be P=[p1,p2,…,pT]P=[p_1,p_2,…,p_T] and the per-step reward sequence of module mkm_k be R=[r1,r2,…,rT]R=[r_1,r_2,…,r_T], where T is the number of training steps. Smoothing. To suppress noise, we apply moving-average smoothing with window size wsw_s to both sequences: Pt(s)=1ws∑i=0ws−1pt+i,Rt(s)=1ws∑i=0ws−1rt+i,P_t^(s)= 1w_s _i=0^w_s-1p_t+i, R_t^(s)= 1w_s _i=0^w_s-1r_t+i, (2) for t=1,2,…,T−ws+1t=1,2,…,T-w_s+1. Slow difference. To capture trends, we compute wdw_d-step differences of the smoothed sequences: ΔPt=Pt+wd(s)−Pt(s),ΔRt=Rt+wd(s)−Rt(s), P_t=P_t+w_d^(s)-P_t^(s), R_t=R_t+w_d^(s)-R_t^(s), (3) for t=1,2,…,T−ws−wd+1t=1,2,…,T-w_s-w_d+1. Lagged correlation. To account for delayed effects of rewards on performance, we compute Pearson correlation at multiple lag steps ℓ=0,1,…,L =0,1,…,L: corrℓ(X,Y)=Cov(Xℓ+1:T′,Y1:T′−ℓ)σ(Xℓ+1:T′)⋅σ(Y1:T′−ℓ),corr_ (X,Y)= Cov (X_ +1:T ,\;Y_1:T - )σ (X_ +1:T )·σ (Y_1:T - ), (4) where (X,Y)(X,Y) can be either the smoothed pair (P(s),R(s))(P^(s),R^(s)) or the differenced pair (ΔP,ΔR)( P, R), T′T is the sequence length after smoothing/differencing, Cov(⋅,⋅)Cov(·,·) denotes covariance, and σ(⋅)σ(·) denotes standard deviation. The raw and differential correlation credits for module mkm_k are: ckraw c_k^raw =max0≤ℓ≤Lcorrℓ(P(s),R(s)), = _0≤ ≤ Lcorr_ (P^(s),\;R^(s) ), (5) ckdiff c_k^diff =max0≤ℓ≤Lcorrℓ(ΔP,ΔR). = _0≤ ≤ Lcorr_ ( P,\; R ). The final correlation credit is a weighted combination: ckcorr=wraw⋅ckraw+(1−wraw)⋅ckdiff,c_k^corr=w_raw· c_k^raw+(1-w_raw)· c_k^diff, (6) where wraw∈[0,1]w_raw∈[0,1] balances the contribution of raw versus differential correlation. A.1.3 EMA Smoothing To prevent credit values from oscillating across iterations, both LLM and correlation credits are smoothed via exponential moving average with update rate α: c~kLLM c_k^\,LLM ←α⋅ckLLM+(1−α)⋅c~kLLM, ←α· c_k^LLM+(1-α)· c_k^\,LLM, (7) c~kcorr c_k^\,corr ←α⋅ckcorr+(1−α)⋅c~kcorr. ←α· c_k^corr+(1-α)· c_k^\,corr. Both credits are initialized to 0. Throughout, ckc_k denotes the raw credit from the current iteration and c~k c_k denotes its exponentially smoothed estimate. A larger α makes the credit more responsive to current observations; a smaller α yields smoother evolution. A.2. Credit Normalization and Fusion Before fusion, the two smoothed credits c~iLLM c_i^LLM and c~icorr c_i^corr are normalized to comparable ranges. Let N be the number of modules in the current pool variant. LLM credit is non-negative and normalized linearly: sillm=c~iLLM∑j=1Nc~jLLM,s_i^llm= c_i^LLM _j=1^N c_j^LLM, (8) with a uniform distribution 1/N1/N applied if all c~jLLM=0 c_j^LLM=0. Correlation credit c~icorr∈[−1,1] c_i^corr∈[-1,1] is normalized via softmax with temperature τ, subtracting the maximum for numerical stability: sicorr=exp((c~icorr−maxj(c~jcorr))/τ)∑j=1Nexp((c~jcorr−maxj(c~jcorr))/τ).s_i^corr= \! (( c_i^corr- _j( c_j^corr))\,/\,τ ) _j=1^N \! (( c_j^corr- _j( c_j^corr))\,/\,τ ). (9) Fusion. The normalized scores are fused into a composite score: si=wLLM⋅sillm+(1−wLLM)⋅sicorr,s_i=w_LLM· s_i^llm+(1-w_LLM)· s_i^corr, (10) where wLLM∈[0,1]w_LLM∈[0,1] controls the relative weight of LLM judgment versus empirical evidence. A.3. Module Selection with UCB To balance exploitation and exploration, an upper confidence bound (UCB) bonus is added to the composite score: ui=si+γ⋅ln(U+1)ui+1,u_i=s_i+γ· (U+1)u_i+1, (11) where sis_i is the fused composite score, U is the total number of times any module from the current pool variant has been selected across iterations, uiu_i is the selection count of module mim_i, and γ is the UCB exploration coefficient. The top-K modules ranked by uiu_i are selected, and their composite scores sis_i are used directly as the weights wiw_i in the reward combination (Eq. 1 of the main paper). Appendix B Hyperparameters Table 4 lists all hyperparameters used in MLREF and the baseline methods. Parameter Value MLREF — Optimization Max iterations I 9 Parallel samples per iteration S 3 Pipeline runs 3 RL evaluation seeds 5 MLREF — Pool Merge Soft update threshold t 0.1 MLREF — Correlation Credit Smoothing window size wsw_s 30 Differential window size wdw_d 30 Max lag steps L 10 Raw correlation weight wraww_raw 0.5 MLREF — Credit Fusion EMA update rate α 0.7 Softmax temperature τ 0.25 LLM credit weight wLLMw_LLM 0.5 MLREF — Module Selection UCB exploration rate γ 0.2 Max modules selected K 5 RL Training Algorithm PPO PPO hyperparameters Same as EUREKA 20 Environment epochs (iterative optimization) 3,000 Environment epochs (final evaluation) 20,000 LLM Configuration Model DeepSeek-V4-Flash Mode Reasoning (thinking mode) Baselines (following original configurations) EUREKA 20 — Iterations 5 EUREKA 20 — Samples per iteration 16 RF-Agent 9 — MCTS updates 80 Table 4: Hyperparameters for MLREF and baseline methods. Appendix C Iterative Evolution Curves This section presents the performance evolution across optimization iterations for all 17 tasks and three domain-level averages. For per-task curve, solid lines denote the best sample for each iteration; for average curves, solid lines and shaded areas represent the mean and standard deviation over tasks. Per-task curves are min-max normalized before averaging to compute domain-level curves. Locomotion. Figure 3 shows evolution curves for all 7 locomotion tasks. Figure 3: Evolution curves for all 7 locomotion tasks. Manipulation. Figure 4 shows evolution curves for all 10 manipulation tasks. Figure 4: Evolution curves for all 10 manipulation tasks. Domain-level averages. Figure 5 reports the average evolution curves. Figure 5: Average evolution curves. Left: all 17 tasks. Middle: locomotion (7 tasks). Right: manipulation (10 tasks). Appendix D LLM Prompt Templates This section presents the prompts used to invoke the LLM at key stages of MLREF. Content enclosed in curly braces denotes dynamic fields populated at runtime. Unless otherwise noted, all prompts share the following system-level preamble. D.1. System Prompt ⬇ You are a reward engineer trying to write reward functions to solve reinforcement learning tasks as effective as possible. Your goal is to write a reward function for the environment that will help the agent learn the task described in text. Your reward function should use useful variables from the environment as inputs. As an example, the reward function signature can be: task_reward_signature_string Since the reward function will be decorated with @torch.jit.script, please make sure that the code is compatible with TorchScript (e.g., use torch tensor instead of numpy array). Make sure any new tensor or variable you introduce is on the same device as the input tensors. D.2. JSON Output Format Constraint All prompts that request structured output from the LLM append the following format specification. Individual prompts below omit this block for brevity. ⬇ You are allowed to reason internally before producing the final answer. Your final output must be a single JSON object that strictly follows the schema below: schema Do not include your reasoning or any extra text. D.3. Initial Reflection Prompts These prompts are issued once before the first iteration to analyze the task and environment. D.3.1 Task Reflection ⬇ Before designing reward functions, briefly reflect on the task to understand the goal. Focus on the following: 1. What is the final success condition of the task? Describe only the end result that indicates success. 2. What are some possible behaviors (including non-intuitive or unexpected ones) that could lead to success? 3. What are some common failure patterns where the agent appears to make progress but never actually completes the task? Important guidelines: - Do NOT assume the task must be solved in well-defined stages or steps. - Do NOT assume the behavior needs to be smooth, stable, or human-like. - Success may arise from aggressive, unstable, or surprising interactions. - Avoid over-structuring the problem; keep the reasoning flexible and open-ended. Keep the reflection concise and avoid over-decomposition. Do NOT design any reward code at this stage. The refined reward code will be designed in the next stage based on your reflections here. D.3.2 Environment Reflection ⬇ Before designing the reward functions, let’s reflect on the following question to better understand the environment: Question: What are the relevant state variables available in the environment that can be used to design the reward functions? Figure out their types and usages. Also identify their shapes if they are tensors. Note that the available variables are those defined in the class with prefix ‘self.‘ in the environment code. Do not use the input variables of the functions, as they may be name aliases and not directly accessible in the reward code. Please include as many variables as possible, since the reward code will be based on these variables. Do NOT design any reward code at this stage. The refined reward code will be designed in the next stage based on your reflections here. D.4. Pool Initialization Prompts These prompts construct the initial module pool from scratch. Generation proceeds in two steps: first specifications, then implementations. D.4.1 Specification Generation ⬇ Design some reward modules that will be included in the module pool. Just provide the function specifications and a brief description of what each module does. Each module should may named as "<aspect>_reward" for clarity. Generate 4 to 6 different reward modules. Tips for designing modules: (1) Consider various reward shaping techniques, such as distance-based rewards, progress-based rewards, and task-specific rewards. (2) Consider make use of every relevant state information available in the environment to design informative reward modules. (3) Ensure that the modules are diverse and capture different aspects of the task. D.4.2 Module Implementation ⬇ Implement the python code for a reward module based on this specification: specification Rules for writing the reward module code: (1) The module should only contain a single function that computes the reward. Do not include any helper functions. (2) The return type of the function must be torch.Tensor. (3) The function name and input variables MUST match the specification. Explicitly specify the type of each input variable and the return type of the function. (4) The code output should be formatted as a python code string: "“‘python … “‘". (5) The code will be run using TorchScript, so it should be compatible with TorchScript. Some helpful tips for writing the reward function code: (1) You may find it helpful to normalize the reward to a fixed range by applying transformations like torch.exp to the reward (2) If you choose to transform a reward component, then you must also introduce a temperature parameter inside the transformation function; this parameter must be a named variable in the reward function and it must not be an input variable. Each transformed reward component should have its own temperature variable D.5. Pool Improvement Prompt From the second iteration onward, this prompt refines the existing module pool based on training statistics and feedback reflection. ⬇ Based on the training statistics of the RL training, please fix the bugs in the existing modules or suggest improvements to the current module pool to better suit the task requirements. The details of the current module pool are as follows: module_pool_details You can select from the following types of actions to improve the module pool: (1) MODIFY: Making small changes to existing modules to better align with the task requirements. You should provide a natural language description of the required changes. The module specification CANNOT be changed in this way. (2) REWRITE: Making big refactors to existing modules to better align with the task requirements. The module specification CAN be changed in this way, but the name of the module should remain the same for easier tracking. (3) ADD: Adding new modules to the pool. You should provide a full module specification for the new module without implementing the code. (4) REMOVE: Removing modules that are redundant or not useful for the task. For each action, please provide a clear natural language explanation of the reasoning behind it. You can refer to the training statistics if applicable. Some helpful tips for analyzing the policy feedback: (1) If the task score is always near zero, then you must consider (a) Removing modules that may mislead the agent. (b) Adding new modules that better capture important aspects of the task. (2) If the values for a certain reward component are near identical throughout, then this means RL is not able to optimize this component as it is written. You may consider (a) Changing its scale or the value of its temperature parameter. (b) Removing it from the pool. (3) If some reward components’ magnitude is significantly larger, then you must re-scale its value to a proper range. (4) The ‘compute_reward‘ function is automatically generated based on the modules and the assembly plan, so you should not directly modify the code of ‘compute_reward‘. Instead, you should modify the modules in the pool and let the system automatically generate the new ‘compute_reward‘ function based on the updated modules and assembly plan. D.6. Feedback Reflection Prompts The feedback reflection prompt varies depending on the outcome of the previous RL training round. We present the three representative variants below. D.6.1 Training Failed ⬇ The refined reward function in the CURRENT iteration did not run successfully, so we roll back to the previous BEST module pool design as the basis of our reflection and improvement. Here is the previous BEST module pool for your reference: module_pool_details Here is the module pool update plan and the error message in the CURRENT iteration: improve_plan training_signal env_plugin Before designing the new reward functions based on the BEST module pool again, let’s reflect on the following questions to better understand the error and identify potential issues in the reward function design: 1. What is the critical error message in the training signal? 2. Where is the error occurring in the code? 3. Based on the provided environment, what could be the potential reasons for this error? 4. How can we avoid this error in the next iteration of module pool improvement? Please think step by step and provide your reflections on these questions. Just answer the questions one by one using natural language, and feel free to provide any additional insights or observations that may be relevant. Do NOT design any reward code at this stage. The refined reward function will be designed in the next stage based on your reflections here. D.6.2 Training Succeeded but Did Not Surpass Best ⬇ The performance of the CURRENT module pool design has regressed compared to the previously BEST performance, so we roll back to the BEST module pool as the basis of our reflection and improvement. Here is the BEST module pool for your reference: module_pool_details Here is the BEST module assembly along with the training statistics from the BEST iteration: best_module_usage_list best_training_signal Here is the module pool update plan in the CURRENT iteration: improve_plan Here is the module assembly and statistics from the CURRENT iteration for comparison: module_usage_list training_signal Before designing new improvement plan based on the BEST module pool again, let’s reflect on the following questions to better understand the statistics and identify potential issues in the current module pool design: 1. What are the key trends and overall performance observed in the training statistics? 2. What are the differences in the module design and assembly between the current and best iterations? Are they responsible for the observed performance regression? 3. How does each reward module contribute to the overall performance? Are there any modules that seem to be more effective or less effective based on the statistics? Provide some suggestions on the future use of these modules. 4. Are there any potential issues or limitations in the current reward function design? How can we address these issues in the next iteration? Please think step by step and provide your reflections on these questions. Just answer the questions one by one using natural language, and feel free to provide any additional insights or observations that may be relevant. Notice that we have rolled back to the BEST module pool, so the next step improvement plan will be based on the BEST module pool. Try to come up with some new ideas to avoid the repeated regression. Do NOT generate any specific plan at this stage. The pool improvement plan will be designed in the next stage based on your reflections here. D.6.3 Training Succeeded and Surpassed Best ⬇ Here is the current module pool details for your reference: module_pool_details Here is the best module assembly along with the training statistics from the CURRENT iteration: module_usage_list training_signal Here is the best module assembly and statistics from the PREVIOUS iteration for comparison: module_usage_list_prev training_signal_prev Here is the module pool update plan in the PREVIOUS iteration: improve_plan Before designing new improvement plan for the module pool, let’s reflect on the following questions to better understand the statistics and identify potential issues in the current module pool design: 1. What are the key trends and overall performance observed in the training statistics? Are there any significant improvements or regressions compared to the previous iteration? 2. What are the differences in the module design and assembly between the current and previous iterations? Are they responsible for the performance regression? 3. How does each reward module contribute to the overall performance? Are there any modules that seem to be more effective or less effective based on the statistics? Provide some suggestions on the future use of these modules. 4. Are there any potential issues or limitations in the current reward function design? How can we address these issues in the next iteration? Please think step by step and provide your reflections on these questions. Just answer the questions one by one using natural language, and feel free to provide any additional insights or observations that may be relevant. Do NOT generate any specific plan at this stage. The pool improvement plan will be designed in the next stage based on your reflections here. D.7. Weight Selection Prompt This prompt asks the LLM to select modules and assign weights, producing the LLM credit used in hybrid weight optimization. ⬇ Here is the modified module pool, along with the specifications of each module. module_pool You should construct the reward function as a linear combination of the modules in the modified pool. Please choose the most appropriate modules and decide their weights to construct the reward function for the given task. You can also refer to the training statistics of the previous RL training and the improvement plan to help you find the most promising modules. Tips for selecting modules: Consider including these modules in the final reward function: (1) The modules whose reward values changed a lot during training, as they are likely to be more effective for RL optimization. (2) The modules that were modified in the improvement plan, as they are expected to better align with the task requirements. (3) The newly added modules in the improvement plan, as they may capture important aspects of the task that were previously missing. Tips for choosing weights: (1) The weights should be non-negative, as all modules are designed to provide positive feedback for desirable behaviors. (2) The weights should have balanced magnitudes to ensure that no single module dominates the reward signal. (3) Consider the relative importance of each module in achieving the task objectives when assigning weights.