Controlling Large Language Model with Latent Action

Adapting Large Language Models (LLMs) to downstream tasks using Reinforcement Learning (RL) has proven to be an effective approach. However, LLMs do not inherently define the structure of an agent for RL training, particularly in terms of specifying the action space. This paper studies learning a compact latent action space to enhance the controllability and exploration of RL for LLMs. Inspired by reinforcement learning from observations, we propose Controlling Large Language Models with Latent Actions (CoLA), a framework that integrates a latent action space into pre-trained LLMs. CoLA employs an inverse dynamics model to extract latent actions conditioned on future tokens, ensuring that the next token prediction is partially influenced by these actions. Simultaneously, CoLA fine-tunes the pre-trained LLM to function as a language world model, capable of incorporating latent actions as inputs. Additionally, CoLA trains a policy model to generate actions within this language world model. The policy model can be trained via behavior cloning to mimic a standard language model or through RL to maximize taskspecific rewards. In this work, we apply CoLA to the Llama-3.1-8B model. Our experiments demonstrate that, compared to RL with tokenlevel actions, CoLA's latent actions enable greater semantic diversity. For enhancing downstream tasks, we show that CoLA with RL achieves a score of 42.4 on the math500 benchmark, surpassing the baseline score of 38.2, and reaches 68.2 when augmented with a Monte Carlo Tree Search variant. Furthermore, CoLA with RL consistently improves performance on agent-based