Discovering Reinforcement Learning Interfaces with Large Language Models

Why You Can't Just Optimize Rewards in Isolation

Setting up a robot to learn a task usually requires humans to manually design what it sees (the observations) and how it gets points (the rewards). This manual engineering is a massive bottleneck in reinforcement learning (RL) deployment. Current state-of-the-art approaches use large language models (LLMs) to automate reward design. However, they almost always assume the observation interface—the structured data the agent receives—is already fixed and tuned.

This assumption breaks in complex environments. If the raw sensor data is poorly structured or lacks critical relational information, even a perfect reward function won't help. The researchers behind LIMEN argue that you cannot treat these as separate problems. They find that observation and reward design are deeply coupled. Optimizing one while leaving the other static leads to catastrophic failure in several domains.

The Problem

The status quo in RL interface design is a bifurcated workflow. Engineers spend weeks hand-coding observation wrappers to extract meaningful features. They then separately design reward functions to shape the learning signal. Recent automation efforts, such as Eureka or Text2Reward, focus heavily on the latter. They use LLMs to write reward code based on natural language descriptions.

However, these methods assume the agent's input is already sufficient. As shown in the evaluation environments in, this is rarely the case.

In the XLand-MiniGrid tasks, the default observations are flattened grids. These lack explicit relational structure (the ability to describe how objects relate to one another). If you evolve only the reward function here, the agent fails. It cannot "see" the relationship between objects, regardless of the reward. Conversely, in continuous control tasks like Panda tracking, observations might be sufficient. But the reward signal is often too sparse (lacking frequent feedback) to drive learning. The RL interface is a bipartite contract. The observation $\phi$ defines the perceived state, and the reward $R$ defines the goal. You cannot reliably optimize one without accounting for the other.

How It Works

LIMEN (Learning Interfaces via MDP-guided EvolutioN) treats interface discovery as a bilevel optimization problem. Instead of searching for a policy (the agent's behavior), it searches for the code that defines the Markov Decision Process (MDP) the policy inhabits. The framework operates through an LLM-guided evolutionary loop, as illustrated in .

The architecture relies on three core pillars:

1. Programmatic Representation: Both the observation mapping ($\phi$) and the reward function ($R$) are executable Python programs. These programs operate on the raw simulator state. They must be JAX-compatible (using a library optimized for high-performance numerical computing) to allow for efficient execution.
2. LLM-Guided Mutation: LIMEN uses an LLM (specifically Claude Sonnet 4.6) as a structured mutation operator. The LLM receives the task description and the parent program. It also receives feedback from previous failures, such as error traces and performance metrics. This allows the "mutations" to be semantic code improvements rather than random noise.
3. Quality-Diversity Search (MAP-Elites): To prevent the search from collapsing into one mediocre strategy, the authors use a MAP-Elites archive. This archive organizes candidate interfaces into a 2D grid. It uses behavioral descriptors like observation dimensionality and reward structural complexity (measured by Abstract Syntax Tree node count). This ensures the system explores both compact and high-dimensional interfaces.

The "inner loop" involves training a standard RL agent (like PPO) from scratch. The resulting success rate serves as the fitness signal for the "outer loop" evolutionary search.

Numbers

The authors demonstrate that joint optimization is necessary. In their main results, LIMEN's joint discovery outperforms "Obs-Only," "Reward-Only," and "Sparse" baselines.

Specifically, they report success rates of 99% for Easy and Medium gridworld tasks. They reach 85% for Hard, 67% for Panda tracking, and 55% for Go1 push recovery.

The failure modes of the ablations are telling. In the Medium XLand task, the Reward-Only baseline achieved only 19% success. Meanwhile, the Obs-Only baseline hit 99%. In the Panda tracking task, the Obs-Only approach failed entirely (0% success). This confirms the paper's thesis. Single-component optimization fails catastrophically on at least one domain in their suite.

The computational cost is significant. An evolution run for MuJoCo tasks (like Go1) takes approximately 6–7 hours on a single NVIDIA L4 GPU. LLM API costs range from $3 to $11 per run. For XLand tasks, the runtime is roughly 1–3 GPU hours. The total cost for all 15 experimental runs was approximately $42.10.

What's Missing

While the results are compelling, there are several gaps:

• Privileged Information Dependency: The framework assumes access to structured simulator state variables. Examples include state.agent.position or state.data.qpos. In real-world robotics, you rarely have access to ground-truth joint velocities. Moving from "simulator state" to "noisy sensor streams" is a massive leap.
• Scaling to Vision: The search space is currently restricted to low-dimensional state vectors. Scaling to high-dimensional visual observations (pixels) is a major open challenge. Evolutionary code generation for image processing would face extreme complexity issues.
• Reliability on Complex Tasks: As seen in, the "Hard" task shows significant variance. Some seeds discover strong interfaces. Others stall below 10% success. This suggests the combinatorial explosion of the search space remains a hurdle.

Should You Prototype This

Depends on your environment.

If you work in a high-fidelity simulator with rich, structured state access, LIMEN is worth a prototype. It automates the tedious parts of the RL pipeline. This includes feature engineering and reward shaping. It can save weeks of manual trial and error. The code is available at https://github.com/Lossfunk/LIMEN.

However, if you are applying this to a real-world robot using raw camera feeds, hold off. The reliance on privileged simulator state is a risk. The "discovered" features might be impossible to sense in the wild. Use LIMEN to find the optimal interface in simulation first. Then, use that as a blueprint for your physical deployment.

Figures from the paper