Reinforcement learning improves language agents by choosing actions directly
RLiD: Reinforcement Learning in the Decision Plane for Language Agents and LLMs
Artificial Intelligence
Summary
Large language models often learn by rewarding sequences of words, which can be inefficient and error-prone. The authors propose a new method called RLiD that skips generating full text and instead chooses from a set of possible actions directly. This approach reduces errors, speeds up learning, and improves performance on tasks like games and puzzles. They also show how to teach regular language models to mimic these better decision-making policies.
What this means in practice
- •For ai product developers: Train language-based agents to perform tasks more efficiently by selecting actions directly instead of generating full text step-by-step.$Commercial implications: This enables faster, more reliable AI assistants and automation tools that perform complex tasks with lower compute costs.
- •For gaming ai engineers: Build game-playing AI agents that achieve higher success rates and lower compute costs by applying direct action choice models to environment interaction.
Authors
Ziyi Chen, Yan Zhang, Jianhui Wei, Jianing Yin, Daoan Zhang
Abstract
Reinforcement learning (RL) has improved the capabilities of large language models (LLMs) across a wide range of tasks. However, token-level RL assigns an action's advantage to the entire generated response, including tokens specifying the parsed action, reasoning chains, and malformed content. Although the environment executes only the parsed action, variation in how the action is expressed in natural language can introduce additional variance into policy-gradient estimates. We propose RLiD (Reinforcement Learning in the Decision plane), a framework that directly optimizes the probabilities of executable action options. RLiD compiles observations, an optional compact memory, and a labeled action menu into the model input, and uses a large decision model (LDM) to select an action in a single forward pass. By eliminating autoregressive token generation and bypassing action parsing, RLiD avoids invalid action outputs and improves RL rollout efficiency. Our analysis shows that, within a fixed text policy, marginalizing over alternative realizations of the same action removes a zero-mean gradient component without changing the expected update. We further distill trained LDM policies into generative LLMs through supervision on their action trajectories. Experiments on ARC, Countdown, and Sokoban demonstrate improved performance or lower computational cost across several settings. On ALFWorld, the complete training and distillation pipeline reaches 99.3% success in 12.7 H100 GPU-hours, compared with 86.6% in 17.3 GPU-hours for token-level RL of the same LLM. With compiled memory, our strongest Qwen2.5-1.5B LDM and distilled generative LLM achieve 100% success on all 134 unseen and 140 seen ALFWorld evaluation games.