Expanding the Action Space of LLMs to Reason Beyond Language

TL;DR

This work introduces ExpA, an Expanded Action Space that decouples language reasoning from external-environment interactions in LLMs, enabling direct routing to tools and environment-specific actions. To train agents in this expanded space, the authors propose EARL, a reinforcement learning framework that uses counterfactual rollouts to encourage exploration of rarely-invoked actions and environments, combined with PPO-style updates. Empirical results on Calc-Bench and Sorting demonstrate that ExpA and EARL outperform vocabulary-restricted baselines, achieving significant improvements in multi-turn and contingent-planning tasks, including perfect Sort-4 accuracy and competitive Sort-5 performance. The approach eliminates the need for external parsers, supports end-to-end training, and shows strong potential for scalable tool use, mathematical reasoning, and embodied AI applications.

Abstract

Large Language Models (LLMs) are powerful reasoners in natural language, but their actions are typically confined to outputting vocabulary tokens. As a result, interactions with external environments -- such as symbolic operators or simulators -- must be expressed through text in predefined formats, parsed, and routed to external interfaces. This overloads the model's language with both reasoning and control duties, and requires a hand-crafted parser, external to the LLM. To address this, we decouple environment interactions from language by internalizing them in an Expanded Action space (ExpA), beyond the vocabulary. The model starts reasoning in the default language environment, but may trigger routing actions and switch to an external environment at any time. From there, the model can only invoke environment-specific actions, receive feedback from the environment, and potentially route back to language as a result. To promote effective exploration of the expanded action space and new environments, we introduce ExpA Reinforcement Learning (EARL) with counterfactual policy optimization. On tasks requiring multi-turn interactions and contingent planning, EARL outperforms strong baselines with vocabulary-constrained actions. It performs robustly across calculator-based multi-task learning and, in the partially observed sorting problem, achieves perfect Sort-4 accuracy while self-discovering an efficient algorithm competitive with classical designs.

Figures (12)

• Figure 1: The Markov Decision Process (MDP) of LLM interacting with external environments. (a) In existing works, LLM is confined to its vocabulary space $\mathcal{V}$ for both reasoning and interaction with external environments, where the latter requires an external parser to detect special patterns. (b) We decouple environment interactions from language by internalizing them as an Expanded Action space (ExpA) $\mathcal{E}$ beyond vocabulary.
• Figure 2: An example rollout with ExpA. Here, <sw> and <cp> route to the swap and compare environments, respectively. Inside them, agents can choose <A>,<B>,<C> as operands. After two operands are chosen, the step procedure updates the latent state $z$ when necessary, routes back to the language environment, and returns the swap or comparison result as a plain-text observation.
• Figure 3: Example questions and the environment-specific actions in Calc-Bench and Sorting.
• Figure 4: Response token length distribution for correct rollouts in GSM8K$^*$.
• Figure 5: ExpA training reward vs. iteration.