ICML 2025

Controlling Large Language Model with Latent Action
Chengxing Jia 1 Ziniu Li 2 Pengyuan Wang 1 Yi-Chen Li 1 Zhenyu Hou 3 Yuxiao Dong 3 Yang Yu 1 4
Abstract
Adapting Large Language Models (LLMs) to
downstream tasks using Reinforcement Learning
(RL) has proven to be an effective approach. How-
ever, LLMs do not inherently define the structure
of an agent for RL training, particularly in terms
of specifying the action space. This paper stud-
ies learning a compact latent action space to en-
hance 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 in-
verse 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 ac-
tions as inputs. Additionally, CoLA trains a pol-
icy model to generate actions within this language
world model. The policy model can be trained
via behavior cloning to mimic a standard lan-
guage model or through RL to maximize task-
specific rewards. In this work, we apply CoLA
to the Llama-3.1-8B model. Our experiments
demonstrate that, compared to RL with token-
level 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, sur-
passing the baseline score of 38.2, and reaches
68.2 when augmented with a Monte Carlo Tree
Search variant. Furthermore, CoLA with RL con-
sistently improves performance on agent-based
1National Key Laboratory for Novel Software Technol-
ogy, School of Artificial Intelligence, Nanjing University, Nan-
jing, China 2The Chinese University of Hong Kong, Shenzhen,
China 3Tsinghua Universi...