🤖 machine learning

Discovering Reinforcement Learning Interfaces with Large Language Models

This paper introduces LIMEN, a large language model-guided evolutionary framework that automatically synthesizes complete reinforcement learning task interfaces—including both observation mappings and reward functions—from raw simulator states and trajectory-level success metrics, demonstrating that jointly optimizing these components is essential for success across diverse domains.

Original authors: Akshat Singh Jaswal, Ashish Baghel, Paras Chopra

Published 2026-05-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Akshat Singh Jaswal, Ashish Baghel, Paras Chopra

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to walk, pick up a cup, or solve a maze. In the world of Reinforcement Learning (RL), the robot doesn't just "learn" on its own; it needs a teacher to tell it two very specific things:

What to look at: Should it look at the color of the floor? The distance to the wall? The angle of its own knee? (This is the Observation).
How to feel good: When it takes a step, does it get a gold star? A tiny point? Or nothing at all? (This is the Reward).

Usually, human experts have to spend weeks or months manually designing these "look-at" and "feel-good" rules. If they get it wrong, the robot never learns.

This paper introduces a new system called LIMEN (Learning Interfaces via MDP-guided EvolutioN). Think of LIMEN as a creative architect and a strict coach working together, powered by a Large Language Model (LLM), to automatically design the perfect teacher for the robot.

Here is how it works, broken down into simple concepts:

1. The Problem: The "Blind" Robot

Imagine you put a robot in a maze.

The Old Way: You tell the robot, "Look at the raw pixels of the camera" (which is like a blurry, confusing mess of colors) and "Get a point only when you hit the exit." The robot is blind and confused. It might never learn.
The New Way (LIMEN): LIMEN asks an AI (the LLM) to write a computer program that acts as a filter. This program says, "Ignore the blurry colors. Instead, look only at the distance to the nearest wall and the angle of the door." It also writes a new rule: "Get a small point every time you get closer to the door, and a big point when you enter it."

2. The Method: Evolution by Trial and Error

LIMEN doesn't just guess once. It uses a process similar to natural evolution, but instead of evolving animals, it evolves computer code.

The Generation: The LLM writes a bunch of different "teacher" programs (some say "look at the floor," others say "look at the ceiling").
The Test: The robot tries to learn using each teacher.
The Feedback: If the robot fails, the system tells the LLM, "That teacher was bad because the robot couldn't see the door." If the robot does okay, the system says, "Good job, but try making the 'points' system more encouraging."
The Mutation: The LLM takes the best teachers and tweaks them (mutates them) to make them even better. It repeats this cycle 30 times, slowly refining the perfect set of rules.

3. The Big Discovery: You Need Both

The most surprising finding in the paper is that you can't just fix one part of the problem. You have to fix both the "what to look at" and the "how to feel good" at the same time.

The authors tested this with two types of tasks:

The Maze (Gridworld): Here, the robot was failing because it couldn't see the relationships between objects (like "the key is next to the door"). If you only fixed the "points" system but left the "vision" messy, the robot still failed. It needed a better "lens" to see the world.
The Robot Arm (Continuous Control): Here, the robot could see everything perfectly, but it was failing because the "points" were too sparse (it only got a point at the very end). It needed a better "coach" to give it feedback along the way.

The Analogy:

If you try to teach a student to play piano by only giving them better sheet music (Observation) but no feedback on their playing (Reward), they might not improve.
If you give them a great teacher who critiques every note (Reward) but the sheet music is scribbled nonsense (Observation), they still can't learn.
LIMEN realized that to succeed, you need to rewrite the sheet music and hire the perfect teacher simultaneously.

4. The Results

The team tested LIMEN on five different tasks, from simple mazes to complex robot balancing acts.

The Result: When LIMEN designed both the vision and the reward system together, the robots learned to solve the tasks with high success rates (up to 99% on mazes).
The Failure of Half-Measures: When they tried to evolve only the vision or only the reward, the robots failed catastrophically on at least one of the tasks.

Summary

In short, this paper shows that we can stop manually engineering the rules for robots. Instead, we can use an AI to automatically write the code that tells the robot what to see and how to be rewarded. By evolving these two things together, the system discovers clever, human-like strategies (like "look at the distance to the goal" or "give a bonus for staying upright") that make learning much faster and more reliable.

It's like automating the job of the "game designer" for AI, ensuring the game is playable and fair so the AI player can actually win.

Technical Summary: Discovering Reinforcement Learning Interfaces with Large Language Models

Problem Formulation

Reinforcement learning (RL) systems rely on environment interfaces—specifically observation mappings ( $\phi$ ) and reward functions ( $R$ )—to define the Markov Decision Process (MDP) an agent experiences. While recent advances have utilized Large Language Models (LLMs) to automate reward design, these approaches typically assume a fixed, pre-tuned observation space. In many scenarios, raw simulator states contain poorly structured information that hinders learning, and manual engineering of both observations and rewards remains a significant bottleneck.

This paper addresses the problem of RL interface discovery from raw simulator state. The objective is to jointly synthesize an observation mapping $\phi: S \to O$ and a reward function $R: S \times A \times S \to \mathbb{R}$ to maximize a trajectory-level success metric $F$ . The authors formalize this as a bilevel optimization problem where the outer loop searches for the optimal interface $(\phi, R)$ , and the inner loop trains an RL agent $\pi$ on the induced MDP $M_I = (O, A, T_\phi, R)$ to evaluate performance. The search space consists of executable programs operating directly on simulator state variables.

Methodology: LIMEN

The authors propose LIMEN (Learning Interfaces via MDP-guided EvolutioN), an LLM-guided evolutionary framework designed to discover effective interfaces.

Core Components

Executable Program Representation: Both the observation mapping and reward function are represented as Python programs using JAX-compatible numerical operations. This allows the search to alter the dimensionality, structure, and semantics of the agent's input space and reward signal, rather than just learning embeddings within a fixed space.
LLM-Guided Evolution: The framework treats the LLM as a structured mutation operator. Starting from a task description, the system generates candidate interfaces. In subsequent iterations, the LLM receives a prompt containing the task description, a parent interface selected from an archive, and feedback from top-performing and failed candidates (including error traces). This feedback loop steers the LLM away from negative code patterns.
Quality-Diversity Search (MAP-Elites): To prevent the search from collapsing into a single strategy, LIMEN employs a MAP-Elites archive. The archive is structured by two behavioral descriptors:
- Observation Dimensionality: Binned into uniform ranges (e.g., 1–50, 51–100 features).
- Reward Structural Complexity: Measured by the Abstract Syntax Tree (AST) node count.
  This ensures the maintenance of structurally diverse solutions (e.g., compact observations with simple rewards vs. complex observations with intricate shaping).
Evaluation Cascade: To manage computational costs, candidates undergo a "short-budget" cascade filter. Interfaces must exceed a minimal success threshold in a truncated training run before proceeding to full multi-seed evaluation. Fitness is defined as the mean success rate across evaluation seeds.

Experimental Setup

The authors evaluated LIMEN across five novel tasks spanning discrete reasoning and continuous control:

XLand-MiniGrid (Discrete): Three tasks of increasing compositional complexity (Easy: object pickup; Medium: relational placement; Hard: multi-step rule chains). Default observations are flattened egocentric grids.
MuJoCo Robotics (Continuous): Two tasks (Go1 Push Recovery: quadruped balance under force; Panda Tracking: manipulator trajectory following). These tasks feature sparse rewards and high-dimensional control.

The experiments compared Joint Discovery (LIMEN) against three baselines:

Sparse: Raw observations with binary success rewards.
Obs-Only: Evolving observations with a fixed reward.
Reward-Only: Evolving rewards with fixed raw observations (instantiating methods like Eureka).

Key Results

Superiority of Joint Optimization: LIMEN consistently outperformed all baselines. It achieved high success rates across all tasks (99% on Easy/Medium, 85% on Hard, 45% on Panda, 48% on Go1).
Catastrophic Failure of Single-Component Optimization:
- Reward-Only failed on compositional gridworld tasks (19% on Medium, 1% on Hard), indicating that even well-structured rewards cannot compensate for unstructured observations in these domains.
- Observation-Only failed entirely on the Panda tracking task (0% success), demonstrating that raw observations are insufficient when the reward signal is too sparse to guide learning.
- This confirms that different domains fail for fundamentally different reasons, necessitating co-design.
Necessity of Evolutionary Feedback: An independent LLM sampling baseline (30 random samples without evolutionary selection) performed significantly worse than LIMEN (e.g., 2.1% vs. 97% on Medium). This indicates that the LLM's prior over interface designs is insufficient alone; the iterative evaluate-and-refine loop is essential for navigating the combinatorial space.
Robustness: Evolved interfaces retained task-relevant structure under distribution shifts (e.g., increased push forces or tracking speeds), degrading gracefully rather than collapsing, suggesting the discovery of generalizable structural patterns.

Analysis of Discovered Interfaces

The analysis of evolved programs revealed consistent, interpretable patterns resembling manual RL engineering strategies:

Observations: Frequent construction of relative geometric features (normalized distances, directional indicators), multi-scale encodings, and explicit task phase indicators.
Rewards: Incorporation of potential-based shaping, milestone bonuses for phase transitions, and smoothness penalties.
Structural Discovery: In the Go1 task, the evolution discovered a structural change in the reward function (removing an "uprightness gate" on position rewards) that significantly improved learning dynamics, demonstrating the system's ability to modify the structure of the reward, not just its parameters.

Significance and Claims

The paper claims that automatic construction of RL interfaces from raw state can substantially reduce manual engineering effort. The primary contribution is the demonstration that observation and reward components often benefit from co-design, as optimizing them in isolation leads to catastrophic failure in at least one domain within the evaluation suite.

The authors position LIMEN as a step toward automating the "inductive bias" of RL problems. By searching over executable programs that define the induced MDP itself, the method produces interpretable and transferable code artifacts rather than learned neural embeddings. The work highlights that while current methods focus on reward design, the observation interface is equally critical and often the limiting factor in complex, compositional tasks.

Limitations: The authors acknowledge that the approach relies on access to structured simulator state variables (privileged information) and a reliable trajectory-level success metric. Furthermore, the computational cost remains a bottleneck, scaling with RL training requirements, though JAX-based parallelism mitigates this for the current experiments. Future work is suggested for larger-scale environments and vision-based observations.