Controlling Large Language Model with Latent Action

Chengxing Jia, Ziniu Li, Pengyuan Wang, Yi-Chen Li, Zhenyu Hou, Yuxiao Dong, Yang Yu
Proceedings of the 42nd International Conference on Machine Learning, PMLR 267:27331-27372, 2025.

Abstract

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 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, 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 tasks without degrading the pre-trained LLM’s capabilities, unlike the baseline. Finally, CoLA reduces computation time by half in tasks involving enhanced thinking prompts for LLMs via RL. These results highlight CoLA’s potential to advance RL-based adaptation of LLMs for downstream applications. The CoLA model is available at https://huggingface.co/LAMDA-RL/Llama-3.1-CoLA-10B.

Cite this Paper

BibTeX
@InProceedings{pmlr-v267-jia25e, title = {Controlling Large Language Model with Latent Action}, author = {Jia, Chengxing and Li, Ziniu and Wang, Pengyuan and Li, Yi-Chen and Hou, Zhenyu and Dong, Yuxiao and Yu, Yang}, booktitle = {Proceedings of the 42nd International Conference on Machine Learning}, pages = {27331--27372}, year = {2025}, editor = {Singh, Aarti and Fazel, Maryam and Hsu, Daniel and Lacoste-Julien, Simon and Berkenkamp, Felix and Maharaj, Tegan and Wagstaff, Kiri and Zhu, Jerry}, volume = {267}, series = {Proceedings of Machine Learning Research}, month = {13--19 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v267/main/assets/jia25e/jia25e.pdf}, url = {https://proceedings.mlr.press/v267/jia25e.html}, abstract = {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 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, 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 tasks without degrading the pre-trained LLM’s capabilities, unlike the baseline. Finally, CoLA reduces computation time by half in tasks involving enhanced thinking prompts for LLMs via RL. These results highlight CoLA’s potential to advance RL-based adaptation of LLMs for downstream applications. The CoLA model is available at https://huggingface.co/LAMDA-RL/Llama-3.1-CoLA-10B.} }
Endnote
%0 Conference Paper %T Controlling Large Language Model with Latent Action %A Chengxing Jia %A Ziniu Li %A Pengyuan Wang %A Yi-Chen Li %A Zhenyu Hou %A Yuxiao Dong %A Yang Yu %B Proceedings of the 42nd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2025 %E Aarti Singh %E Maryam Fazel %E Daniel Hsu %E Simon Lacoste-Julien %E Felix Berkenkamp %E Tegan Maharaj %E Kiri Wagstaff %E Jerry Zhu %F pmlr-v267-jia25e %I PMLR %P 27331--27372 %U https://proceedings.mlr.press/v267/jia25e.html %V 267 %X 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 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, 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 tasks without degrading the pre-trained LLM’s capabilities, unlike the baseline. Finally, CoLA reduces computation time by half in tasks involving enhanced thinking prompts for LLMs via RL. These results highlight CoLA’s potential to advance RL-based adaptation of LLMs for downstream applications. The CoLA model is available at https://huggingface.co/LAMDA-RL/Llama-3.1-CoLA-10B.
APA
Jia, C., Li, Z., Wang, P., Li, Y., Hou, Z., Dong, Y. & Yu, Y.. (2025). Controlling Large Language Model with Latent Action. Proceedings of the 42nd International Conference on Machine Learning, in Proceedings of Machine Learning Research 267:27331-27372 Available from https://proceedings.mlr.press/v267/jia25e.html.

Related Material