PaperlayerView exact revision on arXiv

Source-linked AI summary

Discovering Reinforcement Learning Interfaces with Large Language Models

Akshat Singh Jaswal, Ashish Baghel, Paras Chopra

arXiv:2605.03408v1cs.LGcs.AI

TL;DR

RL research still requires substantial manual design of observation and reward interfaces, while prior LLM reward-design methods generally keep observations fixed. LIMEN jointly evolves executable observation and reward programs using policy-training feedback, and joint evolution avoids catastrophic failure across the paper’s five tasks. The results indicate that observations and rewards often benefit from co-design, although reliable task-success evaluation and substantial computation remain practical constraints.

  • Problem

    Manual RL interface design remains a bottleneck, and existing LLM reward-design methods generally assume fixed observations rather than discovering complete interfaces.

  • Method

    LIMEN uses LLM-guided evolutionary search over executable observation and reward programs, evaluating candidates through policy learning and a quality-diversity archive.

  • Results

    Joint evolution avoids catastrophic failure across all five tasks, whereas observation-only and reward-only optimization each fail on at least one domain.

  • Takeaways & Limitations

    Observations and rewards often benefit from co-design, with task demands producing different observation- and reward-related learning limitations across domains.

  • Takeaways & Limitations

    LIMEN relies on a reliable external task-success metric, and its evaluation cost grows with RL training, especially for larger or high-dimensional environments.

Abstract

from arXiv · show

Reinforcement learning systems rely on environment interfaces that specify observations and reward functions, yet constructing these interfaces for new tasks often requires substantial manual effort. While recent work has automated reward design using large language models (LLMs), these approaches assume fixed observations and do not address the broader challenge of synthesizing complete task interfaces. We study RL task interface discovery from raw simulator state, where both observation mappings and reward functions must be generated. We propose LIMEN (Code available at https://github.com/Lossfunk/LIMEN), a LLM guided evolutionary framework that produces candidate interfaces as executable programs and iteratively refines them using policy training feedback. Across novel discrete gridworld tasks and continuous control domains spanning locomotion and manipulation, joint evolution of observations and rewards discovers effective interfaces given only a trajectory-level success metric, while optimizing either component alone fails on at least one domain. These results demonstrate that automatic construction of RL interfaces from raw state can substantially reduce manual engineering and that observation and reward components often benefit from co-design, as single-component optimization fails catastrophically on at least one domain in our evaluation suite.

1 Introduction

RL interfaces determine what agents observe and how they are rewarded, but manually designing these components remains a bottleneck. LIMEN jointly evolves executable observation and reward programs using policy-learning feedback, and joint evolution avoids catastrophic failure across five evaluated tasks.

  • Manual observation and reward design remains a critical bottleneck because these choices affect learning efficiency, exploration, and final policy performance.
  • LLM-based reward-design methods generally assume a fixed observation interface, leaving automatic discovery of complete interfaces relatively underexplored.Raw state can be poorly structured, whereas carefully designed observations can simplify learning.
  • LIMEN jointly optimizes observation mappings and reward functions as executable programs through LLM-guided evolutionary search and policy-training evaluation.Candidate interfaces are refined using a quality-diversity archive and a trajectory-level success metric.
  • The evaluation spans novel gridworld reasoning tasks and robotic control environments.The experiments test interface generation across different task types rather than a single environment family.
  • Jointly evolving observations and rewards is the only evaluated approach that avoids catastrophic failure across all five tasks.Observation-only and reward-only optimization each fail on at least one domain, while discovered interfaces exhibit recurring interpretable patterns.

2 Problem Formulation

The paper formalizes an RL interface as an observation mapping and reward function that transform a simulator into an induced learning problem. Interface discovery is posed as bilevel search over executable programs, with trained-policy success providing the outer objective.

  • 2.1 RL Interface and Induced MDP: The simulator is modeled as an MDP with state space S, action space A, transition dynamics T, and initial-state distribution ρ0.
  • 2.1 RL Interface and Induced MDP: A task-specific success metric F evaluates trained policies over complete episodes and serves as the fitness function.
  • 2.1 RL Interface and Induced MDP: An RL interface is the pair (ϕ, R), where ϕ maps simulator states to observations and R specifies rewards over state-action-state transitions.
  • 2.1 RL Interface and Induced MDP: The interface induces the agent’s learning problem by producing observations o_t = ϕ(s_t) and observation-level dynamics derived from T.
  • 2.2 Interface Discovery Objective: Given an interface, an RL algorithm trains a policy on the induced MDP.
  • 2.2 Interface Discovery Objective: Interface discovery seeks an interface that maximizes task performance under the evaluation metric F.
  • 2.2 Interface Discovery Objective: The optimization accounts for stochasticity from policy initialization, environment randomness, and training noise.
  • 2.2 Interface Discovery Objective: The search is bilevel and operates over executable programs that directly manipulate raw simulator-state variables.

3 Related Work

Prior work has automated reward construction, learned representations over fixed inputs, and applied evolutionary or programmatic search in RL. LIMEN differs by jointly searching observation mappings and reward functions directly from raw simulator state.

  • Reward-design research includes inverse reinforcement learning, imitation learning, preference-based methods, RLHF, and LLM-generated reward code.
  • Representation-learning methods learn embeddings over a fixed observation space rather than changing the observation function supplied to the agent.
  • Evolutionary algorithms and LLM-guided mutation have been used for policy search, hyperparameter tuning, program synthesis, and reward evolution.
  • The paper identifies no prior work that jointly searches over observation mappings and reward functions to construct interfaces from raw simulator state.

4 Method

LIMEN discovers RL interfaces by having an LLM evolve executable observation mappings and reward functions, evaluating candidates through policy-training feedback. A MAP-Elites archive preserves structurally diverse interfaces while fitness reflects trained-policy success.

  • Interface synthesis: LIMEN uses LLM-guided evolutionary search to synthesize executable observation mappings and reward functions from task and environment specifications.Each candidate interface contains programs that operate on simulator state to produce agent inputs and scalar rewards.
  • Interface representation: Candidates are represented as Python programs operating on raw simulator state, with observations limited to 512 features and rewards computed from state transitions.Observation programs use JAX-compatible numerical operations, while reward programs may use environment statistics such as cumulative errors or episode progress.
  • Evolutionary loop: The search initializes and mutates interfaces with an LLM, validates programs, trains policies, estimates fitness, and inserts candidates into an archive before returning its best interface.Experiments evaluate each candidate by training PPO agents across three random seeds.
  • Quality-Diversity Archive and Selection: MAP-Elites maintains diversity using observation dimensionality and reward AST node count as behavioral descriptors for archive niches.The archive is a two-dimensional grid whose axes bin observation size and reward structural complexity.
  • Inner-Loop Evaluation and Fitness: Interface fitness is based on a from-scratch trained policy’s performance in the induced MDP, measured as mean success across evaluation seeds.A short-budget cascade can discard candidates below a user-set success threshold before full multi-seed evaluation, and the framework can accommodate auxiliary or domain-specific objectives.

5 Experiments

LIMEN is evaluated across five tasks spanning gridworld reasoning and robotic control, using fixed PPO configurations and evolutionary interface search. Jointly evolving observations and rewards outperforms component-wise alternatives, while evolution and robustness tests show benefits beyond independent sampling and nominal dynamics.

  • Evaluation Tasks: LIMEN evaluates five tasks spanning XLand-MiniGrid reasoning and MuJoCo locomotion and manipulation.The suite includes three gridworld tasks, Go1 push recovery, and Panda tracking.
  • Evolution Protocol: Each evolution run generates candidate interfaces, trains PPO agents for evaluation, and uses cascade filtering for XLand-MiniGrid candidates.Candidates exceeding 1–5% success proceed to full multi-seed training; MuJoCo candidates receive full training directly.
  • Evaluation Caveat: Additional five-seed evolution runs converge reliably on Easy and Medium but show higher variance on Hard.The reported main results use one evolution run per task.
  • Main Results: Joint discovery reaches 99% Easy, 99% Medium, 85% Hard, 45% Panda, and 48% Go1 success, outperforming observation-only, reward-only, and sparse baselines.The best interfaces are retrained from scratch and evaluated over 10 independent seeds under fixed budgets.
  • Main Results: Reward-only search collapses on Medium and Hard, whereas observation-only search fails entirely on Panda, revealing complementary single-component failure modes.Reward-only reaches 19% on Medium and 1% on Hard; observation-only reaches 0% on Panda.
  • Evaluation Caveat: The fitness function uses mean success rate for fair comparison, while unconstrained generation can produce unnecessarily large observation vectors.The authors cite a 174-feature Easy observation as motivation for future dimensionality penalties.
  • Independent Sampling: Evolution substantially outperforms independent LLM sampling across four tasks, indicating that iterative evaluation and refinement is needed beyond the LLM’s prior alone.Independent samples average 2.1% and 0.8% on Medium and Hard versus 97% and 76% with evolution, and 21.5% and 10.9% on Go1 and Panda versus 55% and 67%.
  • Robustness: Under perturbed robotics dynamics, performance degrades continuously rather than collapsing to zero, though sensitivity varies by perturbation type.Doubling Go1 push force reduces success from 50.3% to 17.8%, while increasing Panda tracking radius reduces it from 53.9% to 29.9%.

6 Analysis

LIMEN’s analysis shows that interface bottlenecks differ across domains: gridworlds are primarily observation-limited, while continuous control can be reward-limited. Joint evolution discovers recurring observation and reward structures, including qualitatively different reward programs that alter learning dynamics.

  • Observation vs Reward Bottlenecks: Gridworld tasks are primarily observation-limited, whereas continuous-control tasks can be reward-limited.Reward-only search achieves only 2% success on Medium and 4% on Hard with default observations, while reward-only evolution reaches 71% on Panda.
  • Observation vs Reward Bottlenecks: In XLand-MiniGrid, observation-only evolution reaches 99% success on Medium while reward-only search reaches 2%.The results attribute this gap to missing structured relational information in default observations.
  • Observation vs Reward Bottlenecks: In Panda tracking, observation-only evolution fails entirely, while reward-only evolution achieves 71% success with a four-term reward.The raw simulator observations appear sufficient, but the sparse task reward provides insufficient learning signal.
  • Observation vs Reward Bottlenecks: Across domains, different failure mechanisms motivate jointly optimizing observations and rewards rather than optimizing either component alone.The analysis distinguishes tasks limited by observation design from those limited by reward shaping.
  • Recurring Interface Design Patterns: LIMEN repeatedly discovers relative geometric features, multi-scale encodings, and explicit task-phase representations in observation programs.These recurring motifs include position vectors, normalized distances, and directional indicators between relevant entities.
  • Recurring Interface Design Patterns: Discovered interfaces resemble manually engineered representation and reward-shaping strategies used by RL practitioners.Reward structures include potential-based shaping, milestone bonuses, and smoothness penalties, while evolution can also change reward gating.
  • Case Study: Structural Reward Discovery: Removing uprightness gating in Go1 increases success from 32% to 55% while adding multi-scale and body-frame position features.The later interface supplies continuous recovery gradients even when the agent is partially unstable.

7 Limitations and Future Work

The approach depends on reliable task-success evaluation and incurs increasing computational cost as RL training is used for evaluation. Future work targets larger, more complex environments, more efficient evolution, and reduced reliance on privileged simulator state.

  • Evaluation and Computational Cost: LIMEN requires an external metric that reliably measures true task success to guide interface evolution.Evolution becomes harder when such a metric is unavailable or difficult to specify.
  • Evaluation and Computational Cost: Evaluation cost grows with RL training, making computational expense the primary practical bottleneck.The experiments mitigate this cost with JAX-based environments supporting highly parallel simulation.
  • Scope and Future Work: The current evaluation does not cover larger-scale environments or tasks with high-dimensional observations such as vision.These settings remain an important direction for future work.
  • Scope and Future Work: The evolutionary search uses a deliberately simple single-candidate iteration scheme with a basic MAP-Elites archive.More sophisticated evolutionary strategies, stronger mutation models, or additional iterations could improve search efficiency or candidate quality.
  • Scope and Future Work: The formulation assumes access to structured simulator state variables, which provides privileged information that may be unavailable in real-world settings.The paper notes that such state is commonly exposed during simulation development and used for reward design and debugging.

8 Conclusion

The paper introduces LIMEN for automatically constructing both observation mappings and reward functions from raw simulator state using trajectory-level success. Across gridworld and robotics tasks, it finds domain-specific bottlenecks and shows that joint co-design avoids failures affecting single-component optimization.

  • Conclusion: LIMEN searches over executable observation and reward programs from raw simulator state using policy-learning evaluation.The framework is guided by an LLM and uses a trajectory-level success metric as its fitness signal.
  • Conclusion: Across gridworld and robotics tasks, compositional problems are primarily observation-limited while continuous-control problems are frequently reward-limited.Single-component optimization fails catastrophically on at least one domain.
  • Conclusion: Jointly optimizing observations and rewards can substantially reduce the manual effort required to formulate effective reinforcement-learning problems.The conclusion presents observation-reward co-design as beneficial across the evaluated task types.

A.1.3 Training Budgets

The appendix specifies the training setup and LLM-driven interface-generation prompts. Initial candidates are generated from task descriptions, while later candidates are refined using performance feedback, failure information, archive diversity, and observation dimensionality.

  • A.1.3 Training Budgets: Training budgets are summarized in environment timesteps for the XMinigrid and MuJoCo settings.The supplied passage identifies Table 3 as the training-budget reference.
  • Training Setup: The policy uses recurrent PPO with a GRU memory module and categorical action outputs.The recurrent architecture includes a GRU hidden size of 512 and normalized observations.
  • Training Setup: The policy and value networks use MLP architectures with layers 512, 256, and 128.The same architecture specification is given for both networks.
  • Interface Generation: Candidate interfaces consist of JAX-compatible get_observation and compute_reward functions, with only JAX imports allowed.The agent receives only the observation output and is trained only on the computed reward.
  • Interface Generation: The prompt instructs candidates to include task-relevant information, normalize features, and provide rewards aligned with real task progress.Observation outputs are one-dimensional float32 vectors and rewards are scalar float32 values.
  • Evolutionary Prompting: The first generation is created without a parent program, whereas subsequent generations improve an existing program using success rate, return, observation dimension, and evolutionary feedback.Feedback includes recent failures, high-performing programs, diverse MAP-Elites cells, and training signals such as variance or plateau detection.
  • Evolutionary Prompting: Mutation guidance changes with parent success rate, from structural reset at 0% to robustness and simplification above 90%.Intermediate ranges target missing features, weak reward signals, failure analysis, and targeted refinement.

B.4 Evolution Seed Variance

Across five independent seeds, evolution is reliable on Easy and Medium tasks but substantially less reliable on Hard tasks. Figure 7 tracks each seed’s running best success rate over 30 iterations.

  • Seed protocol: Five independent seeds were evaluated for XLand-MiniGrid evolution, with each seed’s running best success rate tracked over 30 iterations.Faded dots represent individual candidate evaluations.
  • Easy and Medium: Easy-task seeds converged to 97–100% success, indicating reliable search on tractable tasks.
  • Easy and Medium: Medium-task seeds converged to 87–100% success, also indicating reliable search on tractable tasks.
  • Hard: On Hard, two seeds reached 76% and 69% while three stalled below 10%, revealing substantial seed variance.The passage attributes this to the combinatorial difficulty of jointly discovering multi-stage reward shaping and structured relational observations.
  • Future direction: The main-paper results report Seed 42, while longer evolution, improved selection, or stronger base models are proposed for improving complex-task reliability.

C.2.1 Task

The experiments cover gridworld pickup-and-placement tasks, quadruped push recovery, and Panda trajectory tracking. Discovered interfaces combine task-relevant observations with sparse completion rewards and dense shaping signals.

  • Go1 push recovery: The Go1 task requires surviving torso pushes while remaining upright, returning to the origin, and maintaining heading direction.Success requires surviving the full episode with average position error below 10 cm.
  • Panda tracking: The Panda task requires tracking a moving three-dimensional Lissajous target with mean end-effector error below 2 cm over 500 steps.Trajectory frequencies and phases are randomized at episode start.
  • Easy task: The blue-pyramid task achieved 99% success with a 174-dimensional observation encoding pose, target relations, and a local egocentric map.Its reward combines sparse pickup completion with distance, adjacency, and orientation shaping.
  • Medium task: The yellow-pyramid placement task achieved 97% success with a 102-dimensional observation representing object relations and valid placement locations.The interface exposes task phases and decomposes reward into pickup, delivery, placement, and navigation shaping.
  • Other tasks: The remaining gridworld and continuous-control interfaces achieved 76%, 55%, and 67% success with observation dimensions of 147, 98, and 94 respectively.

F.2 Ablation Experiments

The ablation experiments are summarized in Table 15, while Table 16 reports total compute and LLM cost across all experiments. Overall usage was 9.2M tokens and $42.10 over 36.1 hours.

  • Ablation accounting: Table 15 summarizes runtime and LLM usage for the ablation experiments.
  • Overall accounting: Table 16 reports total compute and LLM cost across all experiments.
  • Overall accounting: 9.2M LLM tokens, $42.10 in API cost, and 36.1 hours of wall-clock runtime were consumed across all experiments.The runtime used a single NVIDIA L4 GPU.
Loading 2605.03408v1…