Designing Observation and Action Models for Efficient Reinforcement Learning with LLMs

Large Language Models (LLMs) have emerged as powerful tools for semantic reasoning, enabling the formalization of tasks that traditionally relied on manual human intuition. This capability extends to environment design in reinforcement learning. While prior research predominantly focuses on reward design, the design of observation and action spaces remains relatively underexplored. We propose LLM-based design of Observation and Action Models (LOAM), a framework leveraging LLMs to construct refined agent spaces from raw environments. To mitigate the computational burden of identifying the best candidate model from stochastic LLM outputs, LOAM incorporates a continuous racing mechanism that dynamically allocates resources to prioritize the most promising configurations without additional training overhead. Empirical evaluations on HumanoidBench and Isaac Lab demonstrate that LOAM consistently outperforms handcrafted baselines in both learning speed and asymptotic performance.

Before applying a reinforcement learning (RL) algorithm, one must first define how an agent observes the environment and what actions it can take. These choices form the interface between the agent and the environment, yet they are often manually designed, whereas RL research has mainly focused on improving algorithms under a given interface. Although recent LLM-based approaches have begun to automate this design process, they have mostly focused on reward design. We propose LLM-based design of Observation and Action Models (LOAM), a framework that uses LLMs to design task-relevant observation and action models. Our results show that, even with the same RL algorithm, better observation and action models can substantially improve performance. To address the variability of LLM-generated designs, we introduce LOAM-Race, which generates diverse candidates and efficiently selects stronger designs through competition.