PaperlayerView exact revision on arXiv

Source-linked AI summary

MLREF: Efficient Module Reuse for Reward Design in Reinforcement Learning via Large Language Models

Chenglin Liu, Xun Wang, Ruishuo Chen, Zhuoran Li, Longbo Huang

arXiv:2608.18827v1cs.LGcs.AIcs.CL

TL;DR

Reward design remains a fundamental reinforcement-learning challenge, and function-level reward evolution struggles to reuse effective modules efficiently. MLREF instead evolves a persistent pool of reusable reward modules and outperforms state-of-the-art LLM-based methods by 25.2% in locomotion and 6.6% in manipulation across 17 tasks.

  • Problem

    Function-level reward optimization cannot efficiently reuse effective modules, contributing to performance oscillation in LLM-based reward evolution.

  • Method

    MLREF evolves a persistent pool of reusable reward modules that are accumulated, refined, and recombined through weighted linear combinations and rollback-equipped pool operations.

  • Results

    25.2% average improvement in locomotion and 6.6% in manipulation across 17 tasks, with stable optimization trajectories.

  • Takeaways & Limitations

    MLREF provides more stable optimization dynamics than prior LLM-based reward design methods across locomotion and manipulation tasks.

  • Takeaways & Limitations

    MLREF’s absolute performance remains modest on sparse-reward manipulation tasks such as Block Stack.

Abstract

from arXiv · show

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

Reward design remains a fundamental reinforcement-learning challenge, and LLM-generated rewards can be unreliable when optimized as monolithic programs. MLREF addresses this by evolving a persistent pool of reusable reward modules, achieving stronger and more stable optimization across 17 tasks.

  • Motivation: Reward-function design is fundamental to reinforcement-learning performance, while sparse rewards provide insufficient learning signals and dense rewards require substantial effort.
  • Motivation: LLM-generated rewards can contain hallucinations, syntactic errors, and semantic misalignment with task objectives, limiting their reliability in practical reinforcement learning.
  • Limitations of prior work: Existing LLM-based workflows predominantly optimize monolithic reward programs, making effective module tracking, credit assignment, and semantic recombination difficult.
  • MLREF: MLREF treats a persistent module pool as the optimization object, accumulating successful modules, refining underperforming ones, and reusing proven components to construct reward functions.The framework combines reflection-based refinement, hybrid credit assignment, and merge with rollback.
  • Results: 25.2% in locomotion and 6.6% in manipulation: MLREF outperforms state-of-the-art LLM-based reward-design methods across 17 Isaac Gym and Bi-DexHands tasks.Evolution analysis reports stable optimization trajectories, while ablations verify contributions from individual components.

2. Related Work

Reward design evolved from human-intensive inverse and preference-based RL toward increasingly automated LLM-based optimization, but existing methods remained limited by computational cost or manual intervention. MLREF distinguishes itself by persistently accumulating, refining, and reusing reward modules throughout evolution.

  • Inverse RL and preference-based RL require substantial human effort and generalize poorly across tasks.
  • Early LLM methods either queried black-box generators per timestep, which is computationally prohibitive, or generated code-form rewards without iterative refinement.
  • Later approaches refined rewards using human feedback on visualized behaviors, but their dependence on manual intervention limited scalability.
  • EUREKA established the first fully automated pipeline for iterative reward optimization, followed by methods using AST-based crossover, heuristic pre-screening, MCTS, and preference-based parameter tuning.
  • MLREF is the first framework to maintain a persistent module pool for systematic module-level accumulation, refinement, and reuse throughout reward evolution.

3. Preliminaries

The preliminaries formalize reinforcement learning as policy optimization in a Markov Decision Process and define reward design as selecting task-relevant reward functions evaluated through resulting policy performance.

  • Markov Decision Processes: An MDP is defined by state and action spaces, transition kernel, reward function, discount factor, and initial state distribution.A policy maps states to action distributions and seeks to maximize expected discounted return.
  • Reward Design: Because reward is the sole task-specific signal in the objective, reward-function choice determines the learned behavior.
  • Reward Design: A reward design problem comprises an environment, candidate reward-function space, policy optimization algorithm, and fitness function evaluating task performance.For each candidate reward R, the resulting policy is defined as π_R = A(E, R).

4. Method

MLREF optimizes a persistent module pool rather than monolithic reward functions, assembling temporary rewards as weighted combinations of reusable modules. Its pool evolves through reflection-guided improvement, module-wise merging, and rollback across iterative RL evaluation.

  • Reward Construction: MLREF constructs each reward as a weighted linear combination of modules, with each module targeting a specific task aspect.The combined modules balance multiple reward desiderata.
  • Pool Evolution: The persistent module pool accumulates, refines, and recombines modules across initialization and iterative optimization, while reward functions remain temporary evaluation instances.Effective modules persist beyond the reward function that introduced them, enabling module-level feedback and reuse.
  • Pool Improvement: From the second iteration, MLREF plans and executes Add, Delete, Modify, and Rewrite operations on individual modules.Separating improvement planning from code execution provides finer-grained control over reward optimization.
  • Merge and Rollback: At each iteration, MLREF evaluates S pool variants in parallel, discards variants below the historical best by margin t, and retains the best version of same-named modules.Surviving modules are consolidated through a module-wise merge, while failed attempts inform later reflection.
  • Reflection: MLREF uses initial reflection before initialization and feedback reflection before later improvements to analyze task, environment, training outcomes, module trends, and errors.Reflection encourages diverse hypotheses while separating analysis from concrete plans and code.

5. Experiments

MLREF is evaluated on 17 locomotion and dexterous manipulation tasks against EUREKA and RF-Agent, achieving the best average performance in both categories. Ablations and iteration-tracking results show that module-level management, reflection, and modular optimization contribute to performance and stability.

  • Evaluation setup: 17 tasks span locomotion and dexterous manipulation challenges from Isaac Gym and Bi-DexHands.The experiments address baseline performance, component contributions, and reward-evolution stability.
  • Baseline comparison: 25.2% improvement: MLREF’s locomotion average normalized score is 3.288 over the best baseline.MLREF achieves the highest mean on 4 of 7 locomotion tasks; on Franka Cabinet, it scores 0.997 versus 0.701 for the nearest baseline.
  • Ablation study: 40.9% degradation: replacing DeepSeek-V4-Flash with GPT-4o causes the most severe ablation decline.MLREF nevertheless continues improving over iterations, indicating that pool evolution can accumulate effective modules from weak initialization.
  • Ablation study: 53.3% degradation: disabling reflection produces the most consistent ablation decline, highlighting reflection’s role in stabilizing optimization.The full MLREF framework achieves the best or near-best performance across the evaluated ablation tasks.
  • Optimization stability: MLREF improves consistently from a lower initial point, whereas EUREKA plateaus after early gains and RF-Agent shows intermediate behavior.On Shadow Hand and Catch Abreast, the baselines exhibit severe oscillation, while MLREF maintains st...

6. Conclusion

MLREF optimizes a persistent pool of reusable reward modules rather than monolithic reward functions, using reflection, hybrid credit assignment, and rollback-equipped merging to support stable iterative evolution. Across 17 tasks, it outperforms state-of-the-art LLM-based reward design methods, with reflection and rollback identified as especially important.

  • Framework: MLREF evolves a persistent module pool by accumulating, refining, and reusing reward components across iterations.It replaces function-level optimization with pool-level operations supported by reflection, hybrid credit assignment, and rollback-equipped merging.
  • Results: 25.2% average improvement in locomotion and 6.6% in manipulation are achieved across 17 tasks over state-of-the-art LLM-based reward design methods.Ablation studies identify reflection as particularly critical, while evolution analysis links rollback to stable optimization trajectories.

Appendix … A.3. Module Selection with UCB

The appendix formalizes MLREF’s hybrid weight optimization through credit assignment, normalization and fusion, and UCB-based module selection. It combines LLM judgments with empirical reward-performance relationships, stabilizes credits across iterations, and uses exploration-aware scores to construct reward combinations.

  • A Weight Optimization Details: The appendix organizes hybrid weight optimization into credit assignment, credit normalization and fusion, and module selection.This pipeline provides the mathematical formulation of the strategy discussed in the main paper.
  • A.1. Credit Assignment: MLREF assigns each module both LLM credit and correlation credit.These two scores capture prompted module preferences and relationships between module rewards and training performance.
  • A.1.1 LLM Credit: After pool initialization or improvement, prompted module weights are normalized to sum to 1 as LLM credit, while unselected modules receive zero credit.The resulting credit is defined for the current iteration.
  • A.1.2 Correlation Credit: Correlation credit uses smoothed performance and module-reward sequences, slow differences, and Pearson correlations across multiple lags.The method combines raw and differential correlations, with w_raw controlling their relative contribution.
  • A.1.3 EMA Smoothing: Exponential moving averages stabilize both credit types across iterations.Both credits start at 0; larger α responds more to current observations, whereas smaller α produces smoother evolution.
  • A.2. Credit Normalization and Fusion: Normalized LLM and correlation scores are fused into a composite score, with w_LLM controlling the balance between LLM judgment and empirical evidence.LLM credit is linearly normalized, while correlation credit in [-1, 1] is normalized with a temperature-τ softmax.
  • A.3. Module Selection with UCB: An upper confidence bound bonus balances exploitation and exploration in module selection.The top-K modules ranked by u_i are selected, and their composite scores s_i directly become reward-combination weights w_i.

B. Hyperparameters

This section presents Table 4, which specifies the hyperparameters used by MLREF and the baseline methods. It serves as the paper’s implementation reference for these methods.

  • Table 4 lists the hyperparameters used in MLREF.
  • The table also reports the hyperparameters used by the baseline methods.
  • Together, these entries define the hyperparameter settings for MLREF and its baselines.

C. Iterative Evolution Curves

This section presents performance evolution across optimization iterations for all 17 tasks, including locomotion and manipulation domain-level averages. Per-task curves show each iteration’s best sample, while averaged curves report task-level means and standard deviations after min-max normalization.

  • C. Iterative Evolution Curves: Evolution curves cover all 17 tasks across optimization iterations, with solid lines marking the best sample at each iteration.For average curves, solid lines and shaded areas represent the mean and standard deviation over tasks.
  • Locomotion: Seven locomotion tasks are evaluated through evolution curves shown in Figure 3.Figure 3 reports the evolution curves for all 7 locomotion tasks.
  • Manipulation: Ten manipulation tasks are evaluated through evolution curves shown in Figure 4.Figure 4 reports the evolution curves for all 10 manipulation tasks.
  • Domain-level averages: Domain-level average evolution curves are reported for all 17 tasks, locomotion, and manipulation in Figure 5.Per-task curves are min-max normalized before averaging to compute domain-level curves.

D. LLM Prompt Templates

The prompt templates structure MLREF’s LLM interactions from task and environment reflection through module-pool initialization, improvement, and outcome-dependent rollback. They constrain generation with reusable system and JSON-format instructions while separating reflection from reward-code design.

  • D. LLM Prompt Templates: MLREF uses a shared system preamble for key LLM stages, with runtime fields enclosed in curly braces {}.The templates invoke the LLM at key stages, and dynamic fields are populated at runtime.
  • D.1. System Prompt: The system prompt asks the LLM to write effective task-specific reward functions using useful environment variables and a provided function signature.Generated reward functions are decorated with @torch.jit.script.
  • D.2. JSON Output Format Constraint: Structured-output prompts require a single JSON object matching {schema}, without reasoning or extra text.This format specification is appended to prompts that request structured output.
  • D.3. Initial Reflection Prompts: Initial reflection separately examines the task’s success condition, possible successful behaviors, failure patterns, and relevant environment state variables.The prompts prohibit reward-code design during reflection and emphasize flexible, non-overstructured reasoning.
  • D.3.2. Environment Reflection: Environment reflection identifies accessible self.-prefixed variables, their types, usages, and tensor shapes for subsequent reward design.The prompt requests many relevant variables while excluding function-input aliases and postponing code generation.
  • D.4. Pool Initialization Prompts: Pool initialization first generates 4 to 6 diverse reward-module specifications, then implementations, using shaping techniques and relevant state information.The initialization process proceeds in two steps: specifications followed by implementations.
  • D.5. Pool Improvement Prompt: From the second iteration onward, training statistics and feedback guide module-pool changes through MODIFY, REWRITE, ADD, or REMOVE actions.MODIFY preserves the module specification, whereas REWRITE may change it; each action requires a natural-language rationale.
  • D.6. Feedback Reflection Prompts: When training fails or performance regresses, feedback prompts roll back to the previous BEST module pool and analyze errors or current-versus-best statistics before improvement.The failed-training variant diagnoses the error message, code location, causes, and prevention steps; the regression variant compares best and current assemblies and training signals.

D.4.2 Module Implementation

The module implementation prompt requires Python reward code that follows a strict single-function interface and returns a torch.Tensor. It also provides optional guidance for reward normalization and temperature-controlled transformations.

  • Interface constraints: Reward modules must contain exactly one reward-computing function without helper functions, with specified names, inputs, types, and a torch.Tensor return type.The function signature must match the specification, including explicit input and return annotations.
  • Output format: The generated implementation must be formatted as a Python code string for execution by the TorchS environment.The required format is a fenced python string.
  • Reward transformations: Reward normalization may use transformations such as torch.exp, but every transformed reward component must have its own named temperature parameter that is not an input.Temperature parameters must be introduced inside the reward function.
Loading 2608.18827v1…