Controlling Large Language Model with Latent Action
Authors: Chengxing Jia
Ziniu Li
Pengyuan Wang
Yi-Chen Li
Zhenyu Hou
Yuxiao Dong
Yang Yu
Venue: ICML 2025 poster (2025)
Keywords: Controllable language model
latent action model
Primary area: deep_learning->large_language_models
TL;DR: Develop a new architecture of language model to control language with latent action

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 **Co**ntrolling Large Language Models with **L**atent **A**ctions **CoLA**, a framework that integrates a latent action space into pre-trained LLMs. **CoLA** employs an \emph{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 \emph{language world model}, capable of incorporating latent actions as inputs. Additionally, **CoLA** trains a \emph{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 \emph{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  \url{https://huggingface.co/LAMDA-RL/Llama-3.1-CoLA-10B}.

Controlling Large Language Model with Latent Action
Chengxing Jia
Ziniu Li
Pengyuan Wang' Yi-Chen Li
Zhenyu Hou?
Yuxiao Dong Yang Yu 14
Abstract
tasks without degrading the pre-trained LLM's capabilities, unlike the baseline. Finally, CoLA reduces computation time by half in tasks in- volving enhanced thinking prompts for LLMs via RI. These results highlight CoLA's poten- tial to advance RL-based adaptation of LLMs for downstream applications. The CoLA model is available at https: http:/::///ggggnngggggco co/ LAMD--RL/Ll 3atimma--.11111---- 1-Co-A-10B.
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 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 imputs. Additionally, CoLA trains pol- icy model to generate actions within this language world model. The policy model can be trained via behavior cloning to mimic 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 Ass 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
1. Introduction
Large Language Models (LLMs) (OpenAl et al., 2023; Dubey et al., 2024) exhibit exceptional proficiency in pro- ducing coherent and contextually grounded text, demon- strating state-of-the-art performance across diverse tasks including translation, summarization, and logical reason- ing. Recently, there has been growing interest in adapting L.LSs to downstream tasks through reinforcement learn- ing (RL) (Stiennon et al., 2020; Ouyang et al., 2022). The effectiveness of RL approaches critically depends on a well- crafted formulation of key elements, namely states, actions, rewards, and transitions (Sutton & Barto, 1998). Extensive research has demonstrated that a carefully designed formu- lation only accelerates RL training but also enhances its overall performance upper bound (Pang et al., 2019; Jia et al., 2024). In the context of LLMs, states typically corre- spond to the contextual information available to the model, while rewards are often tailored to specific objectives.
However, the design of actions and transitions remains highly flexible and open optimization, presenting both opportunities and challenges. A common approach to fram- ing actions and transitions involves treating the LLM as an integrated system, employing one-token-one-actee formulation, as seen in works like (Rafailov et al., 2023; Li et al., 2023b; Zhong et al., 2024; Li et al., 2024a; Pang et al., 2024), where each token itself corresponds to an ac- tion. While straightforward, this formulation results in an excessively large action space, exemplified by the 128K- token vocabulary size of Llama-3-series models (Dubey et al., 2024) and the 256K of Gemma-2 models (Team et al., 2024). The expansive action space poses significant chal- lenges in computational efficiency and training feasibility.
National Key Laboratory for Novel Software Technol- ogy. School of Artificial Intelligence, Nanjing University, Nan- jing. China 'The Chinese University of Hong Kong. Shenzhen, China 'Tsinghua University, Beijing, China "Pazhou Laboratory (Huangpu), Guangzhou, China Correspondence to: Yang Yu <yuy many@nju.edu.n>>....
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
To address the above challenges, this paper explores the

---

Controlling Large Language Model with Latent Action
question of how to define a well-structured action space and design effective RL approaches for LLMs.. We draw inspi- ration from the literature on "reinforcement learning from observations only" (Torabi et al., 2019b; Sun et al., 2019; Zhu et al., 2020; Kidambi et al., 2021), a setting where only observations are provided, while actions and the underlying transition dynamics are absent from the dataset-a scenario analogous to the challenges faced in LLMs where only token sequences are available in the dataset, but much structural in- formation is missing and hidden. Extensive research in "RL from observation only" suggests that learning latent actions and transition models significantly enhances controllabil- ity and generalization, as latent actions create a compact representation of the decision space while learned transi- tion models enable prediction of future states from current observations alone, together enabling agents generalize effectively to new scenarios. Building this insight, we aim to construct a framework that reformulates the language model as a transition model augmented with additional in- puts of latent actions. A key advantage of our approach is that the size of the latent action space is substantially smaller than the token-level aitiionvocabularyy size of the LLM. This reduction in dimensionality not only mitigates the computational inefficiencies associated with large action spaces but also has the potential to accelerate RL training and unlock its full effectiveness.
models lacked the ability to output a thinking format, our approach improved prompting efficiency by 2x, enabling the models adopt this format and produce correct answers more effectively. Further, we propose fine-tuning the lan- guage world model under latent action guidance. Compared with standard supervised fine-tuning on the base model, CoLA performs better on multiple tasks, including pref- erence alignment with an average win rate of 64%. 11% improvement on math reasoning. and better performance on two agentic multi-turn tasks, including Alfworld (Shridhar et al., 2020) and Scienceworld (Wang et al., 2022). And CoLA also demonstrates better alignment performance and robustness against reward hacking when the reward model is sub-optimal. Codes for training CoLA will be available at http://ittub.om//LAMMAAAA
2. Preliminaries
Reinforcement Learning in LLMs. We introduce the ba- sic settings of reinforcement learning (RL) (Sutton & Barto, 1998). In RI., problems are often framed by a Markov Deci- sion Process (MDP) (Puterman, 1994), which contains a tu- ple M < A, the set of all contextual information (21, (T), where 2 the token step f and denote the sequence by Fi:t- And R is the reward model of current state. In our paper, we mainly consider an outcome reward model (ORM) R(IIT), which is a sparse reward and only gives the reward signal at the end of generation. AA is the action space containing all the actions a, at each step, which control the transition T to transition from the current state Z1:r to the next state T1:++1 by transition distribution T( 12 at). The goal of RL is to find a policy that selects actions maximize the cumulative reward. For the action, standard LLMs adopt each token I2 as an action to generate and align in RL.
T, R >, In language, the state space is
The technical question now becomes: how to effectively learn this latent action model and transition model, pos- sibly at a low cost? To address this question, we propose Controlling Large Language Models with Latent Actions (CoLA) that augments pre-trained LLM with additional latent actions; see Figure 1. In CoLA, a pre-trained LLM is utilized to provide well-trained representations to expe- dite the training process. Based on the embeddings, we additionally introduce an auxiliary inverse dynamics model to construct the latent action space from token sequences. sequences Then a merge module inserts the extracted latent action into the pre-trained embeddings to complete the transition dynamics, where the observation transitions to the next ob- servation guided by the latent action. Based on the transition, we incorporate a policy for selecting latent actions based on historical context. By learning the latent action policy on a certain reward signal, we achieve a more flexible and controllable language adaptation process.
RL from Observation. When learning data only includes observations Ii:t, necessary to consider how to per- form "Learning from Observation" (LfO). General LIO ap- proaches either use small amount of labeled actions to assist in learning ground truth actions, directly match the distribution of expert data without reward signals. In our setting, we believe there is no suitable way to label ground truth actions in language, and reward signals can be utilized for learning. Therefore, we aim to directly learn latent ac- tions and the underlying transitions T( [Zr,,, a:) and use latent action policy Tt (aae(x1))) for RL.
We conducted experiments to verify the effectiveness of CoLAA Using the 1lama-3.1-88 (Dubey et al., 2024) model as the foundation, we successfully transformed it into la- tent action-controlled model by training it on a large corpus. The corpus is from open-source data. This latent action control enhances the diversity of the generated outputs com- pared to the base model. Then, we compared the efficiency of RL on the trained CoLA and base models. Experiments in the Countdown Game showed that, although all the initial
3. Framework
In this section, we present the framework of CoLA, which aims to construct the language latent action space and the underlying transitions, which we call the language world model, in an unsupervised manner. To be more efficient, we consider converting pre-trained token-level LLM to a
2

---

Controlling Large Language Model with Latent Action
The j - -h token
aj
The j - th latent action
Prompt: Find three examples of British English slang.
Prompt: Find three examples of British English slang.
Xo
X1
X2
X3
X44
X5
X6
X
X
X1
X2
X3
X44
X5
X
X7
Policy (Inference) / Inverse Dynamics (Training)
Language Model
Language World Model
X
X3
X
X5
X6
X7
X8
X1
X2
X33
X4
X5
X
X7
X8
Figure 1. An illustration of latent action control in CoLA. The left the naive decoder-only inference pipeline: and the right the pipeline of CoLA.
latent action model at lower cost. First, describe the design of components in CoLA. Next, we outline how to train CoLA. Finally, we introduce the inference of CoLA.
butions by controlling actions and is the component for RL.
Inverse Dynamics Model finverse: This module takes both the historical state 31 and the next-token Z4 as input, and unsupervisedly extracts discrete latent action ar. For the latent action design, we adopt a codebook C {c J.N of size N, where each c; corresponds to specific latent action. Note that the inverse dynamics model needs future information as input, thus it does not serve as an inference module but only assists training.
We also provide a brief illustration of the latent action model in Figure and compare with the naive decoder-only pipeline.
3.1. Model Design
To realize the idea of latent action control, we seek to unsu- pervisedly extract latent actions from language sequences. When language model accepts input of latent actions, we call it a language world model whose output can be affected by choosing the latent actions. A policy model outputs the latent actions to the language world model. Based on a pre-trained LLM as the base model, we have designed the following modules:
We also show the architecture of CoLA in Appendix A.
3.2. Model Training
After completing the model design, we further introduce how to train these components. First, we should construct the latent action space and underlying world model as the basic decision modules, then we initialize the policy model via action-level behavior cloning. The first two parts require large corpus such a pre-training dataset. After finishing the large-scale training, we conduct latent levell reinforcement learning on the policy model to achieve spe- cific goals or tasks, where we find such a process is much more efficient than previous M-b-edd RI. Finally, due to the limitations of the base model, we also introduce world model fine-tuning methods to accomplish more complex tasks.
Language World Model fworld: This model takes the cur- rent state 11tt and a latent action a, as input, to transition the next-token Ie+1, which corresponds to the underly- ing transition under latent action. Based on pre-trained LLM, we merge the latent action and the embedding of the LLM through a structure with few additional parame- ters, mapping them to the distribution of the next token. The next token distribution should be controllable by the latent actions.
Policy Model π: The policy model takes the current state Ti.t input, and outputs the distribution of latent action dg. Since the language world model is controlled by the latent actions, this module aims to adjust the token distri-
Latent Action Space Learning: Since the latent action unknown, and the base pre-trained LLM cannot be con-
3

---

Controlling Large Language Model with Latent Action
trolled by such unknown conditions, should construct the latent action space and underlying world model from a large corpus in unsupervised manner. We train the inverse dynamics model as an encoder to output latent ac- tion and insert the latent action into the base model, which serves as a conditional decoder. The whole joint train- ing process is like VQ-VAI (van den Oord et al., 2017). However, since VQ- VAE highly prone vocabulary collapse, we employed novel method of direct action assignment to train our model. The details of this method can be found in Appendix A.3.
3.3. Model Inference
To generate each token, CoLA generates a latent action from the policy model and then generates the token from the language world model. Given context Ti:p. we process is:
Step I: ar ~ x Step 2: F(+1 = fworld (F1.t ac).
(1)
Note that compute the next token from the world model greedily. For stochastic generation, we randomly sample actions from the policy model.
Latent Action Policy Behavior Cloning: After construct- ing the latent action space, we initialize the policy model via latent action-level behavior cloning. Specifically, the inverse dynamics model outputs the ground truth latent action, and the policy aims to mimic the latent action label.
4. Experiments
We conduct extensive experiments on benchmarks in mathe- matics, reasoning, and agent tasks. The experimental design is primarily aimed at addressing the following key ques- tions:
Latent Action Reinforcement Learning: Since we have constructed control the latent action level through prior learning, well as separate policy and world models, during the reinforcement learning phase, we directly per- form reinforcement learning at the policy model level for a given reward function. That is, we fix the parameters of the world model, and the policy explores at the latent action level to shift and align the token distribution. We find that, due to the smaller space of latent actions and their more diverse semanties, this approach leads to a more efficient reinforcement learning process.
Can the latent actions effectively enable semantic di- versity in text generation? (Section 4.2)
Can CoLA demonstrate better efficiency over the token- level model in the downstream task fine-tuning stage? (Section 4.3 and Section 4.4)
Can CoLA, with more efficient exploration, mitigate reward hacking? (Section 4.5)
4.1. Experiment Setup
World Model Fine-tuning under Latent Action: Dur- ing our experiments, we found that although models pre-trained directly on large-scale corpora demonstrated high efficiency latent action RI,, the capabilities of the pre-trained models we chose limited our ability to perform more complex tasks, such as preference align- ment, complex mathematical reasoning, and multi-turn agent reinforcement learning. Therefore, we proposed fine-tuning the world model for specific tasks, which we call Tuning under Action Guidance (FTA). By distinguishing the source of the actions responsible for guiding the fine-tuning, we introduced two variants of FTA: FTA from Inverse Dynamics (FTA-I) and FTA from Policy Model (FTA-P). We found that FTA-I is suitable when the fine-tuning data is diverse, while FTA-P is better suited for cases where the fine-tuning data is more limited. Both methods outperformed traditional Supervised Fine- Tuning (SFT) in terms of efficiency. For example, FTA-I can effectively retain the knowledge of the pre-trained model, while FTA-P further enhances fine-tuning perfor- mance. Additionally, further RL built on these methods also demonstrated superior capabilities.
The CoLA model consists of three components: the inverse dynamics model, the world model, and the policy model, with parameter sizes of 1B, 8B, and 2B, respectively. The world model is initialized with Llama-3.1-8B-B
to lever- age existing knowledge as much as possible, Since the se- mantic space of the original LLaMA model has been altered, we conduct continued pre-training on a large-scale dataset to learn the action space and adapt the world model to gen- eration guided by the policy model. We select several open- source datasets, including Slimpajama (Cerebras, 2023), Starcoder (Li et al., 2023a), Proof- (Azerbayev et al., 2024), and WaDao (Yuan et al., 2021), covering general knowledge, code, mathematics, and Chinese and English bilingual content, totaling 1.1T tokens. Due to resource constraints, train only the inverse dynamics model and the world model on 200G randomly selected tokens from this dataset, with 100G of these tokens used for training the policy model to validate the effectiveness of CoLA, More details are provided in Appendix I.I.I.
4.2. The Effectiveness of Latent Actions
For more details of our model training methods, please refer to Appendix B.1.
Inspired by reinforcement learning from observations, we leverage future information to construct a latent action space that should be effective in guiding generation: Does the

---

Controlling Large Language Model with Latent Action
constructed latent action space effectively guide the world model to generate more diverse and higher-quality outputs?
4.3. Efficient Alignment of CoLA in Math Tasks
In section 4.2, we established the validity of the constructed action in math tasks. In this section, aim to further demonstrate that by leveraging actions as guidance, CoLA can effectively facilitate the efficient exploration of LLMs through search methods on downstream tasks.
19
Random Sampling
Lamm 3.1 58 Sampling
Random Taken Sampling
Random Action Sampling
1.7
4.3. 1. THE PERFORMANCE IN MATH REASONING
diversity
We then aim to show that the latent action model can control better in mathematical reasoning. We tune the model on the NuminaMath dataset. We compare training the language world model with policy (FTA-P) with the baseline (Llama- 3. 3.1-8B SFT on the same dataset) in several benchmarks, including math500 (Hendrycks et al., 2021), gsm8k (Cobbe et al., 2021), AIME and Drop (Dua et al., 2019), where the first three are mathematical reasoning tasks, and the fourth is a general reasoning task. Results in Figure 3 (a) show that our model achieves better performance on both math reasoning and general reasoning tasks under math data tun- ing, demonstrating better controllability on reasoning tasks. We also show the pass@ K of math500 between the baseline and CoLA with FTA-P in Figure 3 (b), where our model also shows better searching ability. For RL training, we construct prompts and utilize LLM-specific reinforcement learning methods train policy with 0/1 rule-based re- ward. The prompts are related to MATH and collected from PRM800K (Lightman et al., 2024). After RL, our CoLA model can achieve 42.4 on math500 and outperforms the baseline score of 38.2.
1.6
1.5
1.4
13
Num Training Tokens
Figure 2. The diversity value. The blue line is the diversity of random latent action sampling. The yellow line is the diversity value the base model, and the green one that random token sampling. The red line is the random action sampling diversity scaling from to 10B pre-training tokens.
We aim to evaluate the semantic diversity of our latent action-controlled generation, where the semantic diversity could represent both language diversity and quality. To measure semantic diversity, we introduce text embed- ding model for evaluation. We argue that when the em- bedding similarity between multiple generated contents is sufficiently high, their semantic diversity is low. We chose BGE-M3 (Chen et al., 2024) as the text embedding model and randomly select multiple data prefixes Dval = = from the Vaal as input, generate Na results (1 Na by a certain approach. We define the semantic similarity of the generation as follows:
4.3.2. RESULTS OF MCTS AND MCTS-Q
Our CoLA model, due to the smaller latent action space, reduces the search space, enabling more flexible control. Here, we present action-level MCTS approach. Unlike LLMSs that use step-level or multi-token-level actions, we employ latent actions as the search nodes in MCTS. Due to the significant search time introduced by action-level explo- ration, we propose a Q-uncertainty-batt based pruning MCTS to mitigate the issue of prolonged search times, which we call MCTS-Q: First, we sample a of responses using CoLA math training set and label rewards using the Qwen- Math-2.5-72B (Yang et al., 2024) reward model. Thus, we obtain a dataset of fx. y, a, r. where x is the prompt, y is the response, a is the action sequence, and r is the re- ward. We then train a Q-function using Double DQN (van Hasselt et al., 2016) on this data and use the Bellman er- ror of the Q-function (Sutton & Barto, 1998) to represent uncertainty. More details of MCTS-Q are provided in the Appendix C and Appendix C.3.3. For nodes with low esti- mated uncertainty, which is determined through a threshold, we directly extend the search tree by exploring steps of actions ahead, treating the k+ sstep actions as nodes in the
Na
Σ
Σ
Sim (eir, ,(.)
DDall No (Na 1)
End 1= j=1,jşi
where Sim(;y)) is the cosine-similarity value between two sequences, and we use the reciprocal of the total semantic similarity as the measure of semantic diversity. We evaluate three types of generation: (a) Random action sampling, which randomly samples latent actions for the world model to generate. (b) Base model sampling, which uses the base model to generate. (c) Random token sampling, which randomly samples tokens to generate. From the results in Figure 2, latent action control shows larger semantic diversity, and we also demonstrate that as the number of pre- training tokens increases, the random latent action sampling achieves more diverse generation. We note that the output diversity affects the performance limit of online RL training directly (Li, 2025; Li et al., 2025).
5

---

Controlling Large Language Model with Latent Action
80
78.3
75.0
COLA baseline
CoLA baseline
60
60
52.5
48.4
50
41.0
40
36.0
40
20
30
2.8
0
math500
gsm8k
AIME
Drop
(a) Performance on Reasoning
(b) PASS@ K Math500
Figure 3. Performance of math reasoning. The blue line the CoI.A model, and the yellow line the baseline. (a) Performance on reasoning benchmarks. (b) Performance of pass@K @ on math500.
single-step actions nodes. We compare this search method, which is 68.2 on math500, with three baselines: (1) MCTS using CoLAA model, each node is k-steps action. (2) MCTS using baseline model (Llama-3.1-8B) each node k-steps action. (3) MCTS-Q using baseline model. These baselines on math500 are 65.4, 63.2 and 63.0. We can draw two conclusions: By compar- ing CoLA and baseline, our method achieves better search performance. And by comparing the improvement between MCTS-Q and MCTS, CoLA can better benefit from ex- ploration search tricks, where the baseline cannot achieve improvement from this, implying a large space cannot fit the flcxible method well.
tree search. Otherwise,
shown in Figure 4. We found that neither the base model nor CoLA initially had the ability think in the required
format. However, as shown in Figure 4 (b), our method rapidly developed the ability to respond in the correct for- mat at time step 10, with an efficiency twice that of the base model at time step 20. After emerging with the ability to answer correctly, both models achieved a training prediction accuracy of 10 15%.
However, found that both models struggled to answer correctly. This phenomenon was also mentioned in (Gandhi et al., 2025), attributed to the inherent limitations of the LLaMA model. Therefore, we further fine-tuned the model on specific datasets to validate its performance on more complex tasks.
4.4. The Performance in Agent Tasks
Furthermore, we tested the reinforcement learning efficiency of our model in agent tasks to further validate the effective- ness of our approach.
Alfworld and Scienceworld. We consider more complex RL tasks, specifically those involving multi-turn RL interac- tions. In this category of tasks, towards a specific goal, the language model initially generates a response y = (vi.T from the environment's 's initial prompt z0 = {- Subse- quently, the environment provides feedback x {-|rr based this response, and the language model further gives its replies in light of the environmental feedback and histor- ical interactions. This cycle of interaction continues until the task is either successfully completed or ultimately fails. Corresponding to the outcome of the task, the LLM will receive sparse reward from the multi-turn interactions.
Countdown Game We chose the Countdown Game, in which the LLM provided with a list of numbers and target integer. The L.MM must use each number in the list only once to compute the target through basic arithmetic operations (addition +. subtraction
multiplication x, and division +). Following the approach of DeepSeck- R1, we also designed a format reward. Specifically, the LLM is required to place its reasoning process within '<think>' and "<\think>' lags, and the correct answer within "<answer>' and *< *nwers tags. No addi- tional formatting is allowed. The reward for adhering to the format is 1; otherwise, it is 0. Additionally, we introduced a correctness reward, where a correct answer receives re- ward of 1, and an incorrect answer receives We optimized the base model combined with CoLA. During optimization, the total reward, format reward, and the length of the output
In our experiments, we select two agentic multi-turn turn tasks, including Alfworld and Scienceworid. We begin by fine- tuning the model using dataset (Song et al., 2024) to adapt the model to the corresponding instructions and to output valid actions. For CoLA, we utilize FTA-P. This is followed by online interaction in the RL environment in Alfworld and Scienceworld, both which encompass a multitude of different tasks, such requiring an agent to locate an
6

---

Controlling Large Language Model with Latent Action
CoLA-total rewerd
COLA
format reward
baseline
1.0
baselire-total
400
baseline-format reward
0.8
Time thinking format
300
Reward
Length
200
0.4
Time point of thinking formal
0.2
point right answer
100
10
20
30
40
10
20
30
40
Step
Step
Curves of Reward
(b) Curves of Response Length
Figure 4. Performance of Countdown Game. The blue line is the CoI.A model, and the yellow line is the baseline. (a) Curves of Format
Reward. (b) Curves of Response Length.
standard KL coefficient of 0.01 and another with KL coefficient of 0.00 to explore whether method can more robustly handle reward hacking and align better, since our reinforcement learning process only trains the upper-level latent action policy without altering the underlying language world model. The results in Figure (a) show that our CoLA model can align distinct types of preferences well on stan- dard RLHF (3/4 types of preference), and be more robust against reward hacking (4/4 types of preference). When kl=0.00, we find that it achieves slight advantage over 0.01 in Figure 5 (b), while the baseline completely failed, implying reward hacking of the baseline and that CoLA is more robust to it. We also give an example generated results for KI. 0.00:
Table Performance of CoLA and baseline on Agentic Environ- ments. Seen means the in- distriiibutoon tasks, and Unseen means the out-of-distributtin tasks. Base the baseline model, while CoLA is our model. SFT FTA means tuned model, while RL means model trained by RL. mark improvements tuned model in red and the non-imprwements in blue.
ALFWORLD
SCIENCEWORLD
BENCHMARK
SEEN
UNSEEN
SEEN
UNSEEN
BASE-SFT
68.6
67.9
17.0
17.5
BASE-RL
68.6.e.0
71.6.3.7
18.0.1.0
15.6.1.
CoLA-FTA
75.7
70.9
24.7
20.4
COLA-RL
77.9.2.2
74.6.3.7
28.4.3.7
21.8.1.4
object and bring it to a designated location. More details are provided Appendix C.1.2.
The performance of the initial model and RL model in both tasks is shown in Table 1. Compared to the baseline, fine- tuned model outperforms it by 7.1 score on AlfWorld and 7.7 score on ScienceWorld After RI. training, CoL achieves stable improvements with even greater performance gains and better generalization on unseen tasks.
Instruction: Find the longest river in Africa.
The lengest river in Africa is the Nile River which stretches for about 6, 650 km from its source (Rift Valley) to its delta in Egypt.
CoLA:
4.5. The Advantage of Reducing Reward Hacking
Baseline: As researcher, would like to clarify what you mean by the longest river in Africa." Could you please provide some examples or criteria to help define this term?
Then we turn to the RLHF process. Reward hacking arises from the sub-optimality of reward models. Even if the reward model is imperfect, can we mitigate this issue by optimizing within smaller action space and enabling more efficient exploration? To show the degree of the alignment of certain preference,
evaluate the GPT-4 win rate by Alpaca-Evas Eval (Dubois et al., 2024) on the validation set of each preference data. In standard RLHF, a KL constraint typically introduced, and an excessively small KL con- straint can lead to language capability degradation due to reward hacking. We conduct two KI. experiments: one with
This demonstrates that our approach effectively maintains knowledge and language capabilities, whereas the baseline hacks the reward, leading to degradation into only generat- ing inquiries like "I would like to clarify what you mean".

---

Controlling Large Language Model with Latent Action
100
01
99.8
96.2
96.1
60
KL-0
00
55.8
baseline
90
89.7
50
48
80
40
70
30
60
20
50
10
7.8
6.4
3.5
40
0
0.0
0.0
ACA
BUS
ENT
LIT
AVERAGE
ACA
BUS
ENT
LIT
AVERAGE
(a) Win rate to baseline
(b) Win rate of KL=0 to KI,=0.01 coefficient
Figure 5. GPT-4 win rate in distinet preferences. ACA means academy, BUS means business, ENT means entertainment and LIT means literature. KI. COEF is the KL coefficient. AVERAGE is the average of four tasks. The value larger than 50 means a better alignment. (a) win rate of CoLA relative to baseline. (b) win rate CoLA with KL coefficient 0.00 relative to that with 0.01.
5. Related Works
tions (Torabi et al., 2019a). This makes "learning from ob- servation" highly relevant and challenging problem. Prior works aim to construct latent actions from observation-only data (Sco al., 2022: Baker et al., 2022). For example, learning latent actions from videos to control video and game generation (Liu et al., 2022; Zhang et al., 2024b), These approaches leverage the dynamics between adjacent video frames to model latent actions, which are then used to control diverse content generation (Edwards et al., 2019) and for further RI. agent construction (Schmidt & Jiang, 2024; Ye et al., 2023). This not only enhances controlla- bility (Bruce al., 2024) but also, due to the higher-level nature of latent actions, enables better transferability across different tasks (Liu et al., 2022).
Mme-Token-One-Acee Formulation. Current large lan- guage models (Radford et al., 2019; Brown et al., 2020; Du et al., 2022; Touvron et al., 2023) typically employ transformer architectures (Vaswani et al., 2017) and auto- regressive structure (Radford, 2018) for training and infer- ence, These models directly predict the next token based the historical token sequence. For reinforcement learning in LLMs (Dai et al., 2024; Ouyang et al., 2022; Stiennon et al., 2020), they use individual tokens actions (Rafailov et al., 2023; Li et al., 2024b; Zhong et al., 2024), which we formulation. In this case,
refer to as ame-token-one-actinn the vast token space introduces challenges in exploration and optimization. For exploration, it inefficient to adopt token-level action search methods, often necessitating the use of coarscr-grained process-based search (Zhang al., 2024a). However, it is hard to define the process, and the segmentation of the process often relies on trivial special symbols for segmentation (Lai et al., 2024; Wang et al., 2024). For optimization, the token-level action requires tuning the whole model parameters adjust the token dis- tribution. Due to the poly-semantie nature of parameters in transformers (Ye et al., 2024; Allen-Zhu & Li, 2024), adjusting token distributions for specific task can simulta- neously affect knowledge and language capabilities in other domains, leading to inaccuracy issues (Huang et al., 2023; Liu et al., 2024; Xu et al., 2023) or alignment tax (Guo et al., 2024: Lin et al., 2024: Zheng et al., 2024; Li al., 2025).
6. Conclusion
In this paper, we propose new framework for controllable language adapting. We decompose the language model into bi-level structure, including latent action policy and downstream language generation. We adapt the language model by the guidance of latent policy. Empirical results demonstrate that it exhibits superior performance. Specifi- cally, low-level tuning under high-level guidance can better adapt to specific modes without knowledge degradation, while alignment can more robustly adapt to specific prefer- ences or capabilities. This also motivates us reconsider tuning and alignment, suggesting that should focus more acquiring higher-level patterns rather than fitting spe- cific data. However, our current work still requires broader comparisons due to the limitation of computation resources, such as the effectiveness across multiple base models.
RL with Latent and Compact Actions. In many real- world applications, only observation-only data is available, such as expert videos of robots without corresponding ac-
8

---

Controlling Large Language Model with Latent Action
Acknowledgements
R., Apps, C., Aytar, Y., Bechtle, S., Behbahani, F. M. P., Chan, S. C. Y., Heess, N., Gonzalez, L.., Osindero, S., Ozair, S., Reed, S. E., Zhang, J., Zolna, K., Clune, J., de Freitas, N., Singh, S., and Rocktäschel, T. Genie: Generative interactive environments. In International Conference on Machine Learning, ICML, 2024.
This work supported by National Science Foundation of China (62495093) and Jiangsu Science Foundation (BK20243039). We thank the anonymous reviewers for their helpful suggestions improving the paper. The authors would like to thank Zhipu AI for sponsoring the computa- tion resources used in this work. This work was done while C. Jia interned at Zhipu AI.
Cerebras. Slimpajama- https: httttp:://hgggiggaaee co http///aaeet//eebbrraeeeeeeeeeee jama 627B, 2023.
Impact Statement
Chen, J., Xiao, S., Zhang, P. Luo, K., Lian, D., and Liu, Z. BGE m3-embedding: Multi-lingual, multi-fuuctioonaliyy multi-granularity text embeddings through self- eeff-knoweedge distillation. arXiv preprint arXiv:2402.03216, 2024.
This paper presents work whose goal advance the field of Machine Learning by proposing a controllable language adapting framework. Our approach leverages open- source datasets and models. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M.. Jun, H., Kaiser, L.,, Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10222100022 14168, 2021.
References
Ahmadian, A., Cremer, C., Gallé, M., Fadace, M., Kreutzer, J., Pietquin, O., Üstün, A., and Hooker, S. Back to ba- sies: Revisiting reinforce style optimization for learn- ing from human feedback in llms. arXiv preprint, arXiv:2402. 14740, 2024.
Dai, J., Pan, X., Sun, R., J., Xu, X., Liu, M., Wang, Y., and Yang, Y. Safe RLHF: safe reinforcement learning from human feedback. In International Conference Learning Representations, ICLR, 2024.
Allen-Zhu, Z. and Li, Y. Physics of language models: Part 3.1, knowledge storage and extraction. In International Conference on Machine Learning, ICML, 2024.
Du, Z., Qian, Y., Liu, X., Ding, M., Qiu, J., Yang, Z., and Tang. J. GLM: general language model pretraining with autoregressive blank infilling. In Annual Meeting of the Association for Computational Linguistics, ACL, pp. 320- 335. Association for Computational Linguistics, 2022.
Azerbayev, Z., Schoelkopf, H., Paster, K., Santos, M. D., McAleer, S. M., Jiang, A. Q.. Deng, J., Biderman, S., and Welleck, S. Llemma: An open language model for mathematics. In International Conference on Learning Representations, ICLR, 2024.
Dua, D., Wang, Y., Dasigi, P., Stanovsky, G., Singh, S., and Gardner, M. DROP: A reading comprchension bench- mark requiring discrete reasoning over paragraphs. In Conference of the North Americun Chapter of the Associ- ation for Computational Linguistics: Human Language Technologies, NAACL-HLT, pp.
Baker, B., Akkaya, I., Zhokhov, P., Huizinga, J., Tang, J., Ecoffet, A., Houghton, B., Sampedro, R., and Clune, Video pretraining (VPT): learning to act by watching unlabeled online videos. In Annual Conference on Neural Information Processing Systems 2022, NeurlPS, 2022.
2019.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Vosse A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu,
Dubois, Y., Galambosi, B., Liang, P., and Hashimoto, T. B. Length-controlled alpacaeval: A simple way to debias automatic evaluators. arXiv preprint arXiv:2404.044 2024.
J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Annual Conference Neural Information Processing Systems 2020, NewrIPS, 2020.
Edwards, A. D., Sahni, H., Schroecker, Y., and Jr., C. L. I. Imitating latent policies from observation. In Inter- national Conference on Machine Learning, ICML, vol- une 97 of Proceedings of Machine Learning Research, pp. 1755-1763, 2019.
Bruce, J.. Dennis, M. D., Edwards, A., Parker-Holder J., Shi, Y., Hughes, E., Lai, M., Mavalankar, A., Steigerwald,

---

Controlling Large Language Model with Latent Action
Gandhi, K., Chakravarthy, A., Singh, A., Lile, N., and Goodman, N. D. Cognitive behaviors that enable self- improving reasoners, or, four habits of highly effective stars. arXiv preprint arXiv:2503.0130/ 2025.
O., Gontier, N., Meade, N., Zebaze, A., Yee, M., Uma- pathi, L.. K., Zhu, J., Lipkin, B., Oblokulov, M., Wang, Z. V, R. M., Stillerman, J. T., Patel, S. S., Abulkhanov, D., Zocca, M., Dey, M., Zhang, Z., Fahmy,, N., Bhat- tacharyya, U., Yu, W., Singh, S., Luccioni, S., Villegas, P., Kunakov, M., Zhdanov, F., Romero, M., Lee, T. Timor, N., Ding, J., Schlesinger, C., Schoelkopf, H., Ebert, J., Dao, T., Mishra, M., Gu, A., Robinson, J.. Anderson, C. J., Dolin-Gavit, B.,, Contractor, D., Reddy, S., Fried, D., Bahdanau, D., Jernite, Y., Ferrandis, C. M., Hughes, S., Wolf, T., Guha, A., von Werra, I.., and de Vries, H. Starcoder: may the source be with you! Trans. Mach. Learn. Res., 2023, 2023a.
Guo, Y., Cui, G., Yuan, L., Ding, N., Sun. Z., Sun, B., Chen, H., Xie, R., Zhou, J., Lin, Y., Liu, Z., and Sun, M. Con- trollable preference optimization: Toward controllable multi-objective alignment. In Conference Empirical Methods in Natural Language Processing, EMNLP, pp. 1437-1454, 2024.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, Measuring mathematical problem solving with the MATH dataset. In Proceedings of the Neural Information Processing Systems Track Datasets and Benchmarks 1, NeurIPS Datasets and Benchmarks, 2021.
Li, Y.-C., Zhang, F,, Qiu, W., Yuan, L., Jia, C.,, Zhang, Z., Yu, Y., and An, B. Q-Adapter: Customizing pre-trained Ilms to new preferences with forgetting mitigation. arXiv preprint arXiv:2407.03856. 2024a.
Hong, W., Ding, M., Zheng. W., Liu, X., and Tang. Cogvideo: Large-scale pretraining for text-to-video gen- eration via transformers. In International Conference on Learning Representations, ICLR, 2023.
Li, Z. Can better cold-start strategies improve rl training for Ilms? https//langbbleepooo022330 www.ioosiittooaaaan.oooo Cold-Start Shaut-Strrtegiis--mppoveeeee Training for 688888661686772555228881
2025. Notion
Blog.
Hu, J. Reinforce++: A simple and efficient approach for aligning large language models. arXiv preprint, arXiv:2501 25..02501. 2024.
Li, Z., Xu, T., and Yu, Y. Policy optimization in rlhf: The impact of out-of- of-rrfferennce data. arXiv preprint arXiv:2312. 10584, 2023b.
Huang, L.. Yu, W., Ma, W., Zhong, W., Feng. Z., Wang, H., Chen, Q., Peng, W., Feng, X., Qin, B., and Liu, T. A survey on hallucination in large language models: Prin- ciples, taxonomy, challenges, and open questions. arXiv preprint arXiv:231 .5222,, 2023.
Li, Z., Xu, T., Zhang, Y., Lin, Z., Yu, Y., Sun, R., and Luo, Z. Remax: A simple, effective, and efficient reinforcement learning method for aligning large language models. In International Conference on Machine Learning, ICML, 2024b.
Jia, C., Wang, P., Li, Z., Li, Y., Zhang, Z., Tang, N., and Yu. Y. Bwarea model: Learning world model, inverse dy- namics, and policy for controllable language generation. arXiv preprint ar arXiv:2405.17039, 2024.
Li, Z. Chen, C., Xu, T., Qin, Z., Xiao. J., Luo, Z.-Q.. and Sun, R. Preserving diversity in supervised fine- tuning of large language models. In The Thirteenth International Conference on Learning Representations, 2025. URL https:////// /ttps:://ppeeeeee opwprevieweeeeeeeeeeeeeee. net/ forum? i id=NGEe7B7bSw.
Kalmukov, Y.
Using word clouds for fast identification of papers' subject domain and reviewers' competences. arXiv preprint 1rXiv:2112.14866 2021.
Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let's verify step by step. In International Conference on Learning Representations, ICLR, 2024.
Kidambi, R., Chang, J., and Sun, W. Mobile: Model-based imitation learning from observation alone. Advances in Neural Information Processing Systems, 34:28598- 28611, 2021.
Lin, Y., Lin, H., Xiong, W., Diao, S., Liu, J., Zhang, J., Pan, R., Wang, H., Hu, W., Zhang, H., Dong, H., Pi, R., Zhao, H., Jiang, N., Ji, H., Yao, Y., and Zhang, T. Mitigating the alignment tax of RLHF, In Conference on Empirical Methods Natural Language Processing, EMNLP, pp. 580-606, 2024.
Lai, X., Tian, Z., Chen, Y., Yang, S., Peng, X., and Jia, J. Step-dpo: Step-wise preference optimiza- tion for long-chain reasoning of lmss. arXiv preprint arXiv:2406.1867 2024.
Li, R., Allal, L. B., Zi, Y., Muennighoff, N., Kocetkov, D., Mou, C., Marone, M., Akiki, C., Li, Chim, J., Liu, Q., Zheltonozhskii, E., Zhuo, T. Y., Wang, T., Dehaene, O.,
Liu, M., Zhu, Z., Zhuang, Y., Zhang, W., Hao, J., Yu, Y., and Wang, J. Plan your target and learn your skills: Trans- ferable state-only imitation learning via decoupled policy
Davaadorj, M., Lamy-Poirier, J., Monteiro, J., Shliazhko,
10

---

Controlling Large Language Model with Latent Action
optimization. In International Conference on Machine Learning, ICMI,, volume 162 of Proceedings of Machine Learning Research, pp. 14173-14196 2022.
Shao, Z., Wang, P., Zhu, Q. Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., and Guo, D. Deepseek- math: Pushing the limits of mathematical reasoning in open language models. arXiv preprint, arXiv:2402.03300, 2024.
Liu, X., Khalifa, M., and Wang, L. Litcab: Lightweight language model calibration over short- and long-form responses. In International Conference on Learning Rep- resentations, ICLR, 2024.
Shridhar, M., Yuan, X., Côté, M.-A., Bisk, Y., Trischler, A., and Hausknecht, M. Alfworld: Aligning text and embodied environments for interactive learning. arXiv preprint arXiv:2010.0376. 2020.
OpenAI, Achiam, J., Adler, S., Agarwal, S., Ahmad, L.,, Akkaya, I,, Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.008774 2023.
Song, Y., Yin, D., Yue, X., Huang, J., Li, S., and Lin, B. Y.
Trial and error: Exploration based trajectory optimization for IIm agents. arXiv preprint arXiv:2403.02550 2024.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems 35, pp. 27730-27744, 2022.
Stiennon, N., Ouyang, L.,, Wu, J., Ziegler, D. M., Lowe, R., Voss, C., Radford, A., Amodei, D., and Christiano, Learning summarize with human feedback. In Annual Conference Neural Information Processing Systems 2020, NeurlPS, 2020.
Pang, J., Wang, P., Li, K., Chen, X., Xu, J., Zhang, Z., and Yu, Y. Language model self-improvement by reinforce- The Twelfth Interna- tional Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024, 2024.
Sun, W., Vemula, A., Boots, B., and Bagnell, D. Provably efficient imitation learning from observation alone. In International conference on machine learning, pp. 6036- 6045. PMLR, 2019.
ment learning contemplation.
Pang, Z.-J., Liu, R.-Z., Meng, Z.-Y., Zhang, Y., Yu, Y., and Lu, T. On reinforcement learning for full-length game of starcraft. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4691-4698, 2019.
Sutton, R. S. and Barto, A. G. Reinforcement learning an introduction. Adaptive computation and machine learning. MIT Press, 1998.
Swiechowski, M., Godlewski, K., Sawicki, B., and Mandz- iuk, Monte carlo tree search: a review of recent modifi- cations and applications. Artif. Intell. Rev., 56(3):2497- 2562, 2023.
Puterman, M. I.. Markov Decision Processes: Discrete Stochastic Dynamic Programming. Wiley Series Prob- ability and Statistics. Wiley, 1994.
Radford, A. Improving language understanding by genera- tive pre-training. 2018.
Team, G., Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., Mesnard, T., Shahri- ari, B., Ramé, A., et al. Gemma 2: Improving open language models at a practical size. arXiv preprint mrXiv:248.0000 18, 2024.
Radford, A., WU, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Tian, Y., Yang, S., Zeng, J., Wang, P., Lin, D., Dong, H., and Pang, J. Predictive inverse dynamics models scalable learners for robotic manipulation. arXiv preprint arXiv:2412.15109, 2024.
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly reward model. In Annual Conference on Neural Information Processing Systems 2023, NeurlPS, 2023.
Torabi, F., Warnell, G., and Stone, P. Recent advances in imitation learning from observation. International Joint Conference on Artificial Intelligence, IJCAI, pp. 6325 6331, 2019a.
Schmidt, and Jiang, M. Learning to act without actions. In International Conference on Learning Representations, ICLR. OpenReview.net, 2024.
Torabi, F., Warnell, G., and Stone, P. Recent advances in imitation learning from observation. arXiv preprint 2019b.
Seo, Y., Lee, K., James, S., and Abbeel, P. Reinforcement learning with retion-free pre-training from videos. In In- ternational Conference on Machine Learning, ICML, vol- ume 162 of Proceedings of Machine Learning Research, pp. 19561-19579 2022.
arXiv:
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M., Lacroix, T. Rozière, B., Goyal, N., Hambro, E.,

---

Controlling Large Language Model with Latent Action
Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lam- ple, G. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397) 2023.
Zhang, D., Zhoubian, S., Yue, Y., Dong, Y., and Tang, J. Rest-mets*: LLM self-training via process reward guided tree search. arXiv preprint
2024a.
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Annual Conference on Neural Information Processing Systems, NIPS, pp. 6306-6315, 2017.
Zhang. Z.., Chen, R., Ye, J., Sun, Y., Wang. P., Pang, J., Li, K., Liu, T., Lin, H., Yu, Y., al. Whale: Towards general- izable and scalable world models for embodied decision- making, arXiv preprint arXiv:2411.085
2024b.
van Hasselt, H., Guez, A., and Silver, D. Deep reinforce- ment learning with double q-learning. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 2094- 2100. AAAI Press, 2016.
Zheng, C., Sun, K., Wu, H., Xi, C., and Zhou, X. Balanc- ing enhancement, harmlessness, and general capabilities: Enhancing conversational Ilms with direct RLHF, arXiv preprint arXiv:2403.02551. 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, Attention is all you need. In Annual Conference Neural Infor- mation Processing Systems, NeurIPS, pp. 5998-6008, 2017.
Zhong, H., Feng. G., Xiong, W., Cheng, X., Zhao, L.,, He, D., Bian, J., and Wang. L. Dpo meets ppo: Reinforced token optimization for rlhf. arXiv preprint urXiv:2404.1822. 2024.
Zhu, Z., Lin, K., Dai, B., and Zhou, J. Off-policy imi- tation learning from observations. Advances in neural information processing systems, 33:12402-12413, 2020.
Wang, C., Deng, Y., Lv, Z., Liang, Z., He, J., Yan, S., and An, B. Q": Improving multi-step reasoning for ilms with deliberative planning, arXiv preprini arXiv:2406. 14283, 2024.
Wang, R., Jansen, P., Côté, M.-A., and Ammanabrolu, P. Scienceworld: Is your agent smarter than a 5th grader? arXiv preprint 2tt:::20030054400 2022.
Wettig, A., Gupta, A., Malik, S., and Chen, D. Qurating: Selecting high-quality data for training language models. In International Conference on Machine Learning, ICML, 2024.
Xu, W., Agrawal, S., Briakou, E., Martindale, M. J., and Carpuat, M. Understanding and detecting hallucinations in neural machine translation via model introspection. Trans. Assoc. Comput. Linguistics, 1 11:546-564, 2023.
Yang, A., Zhang, B., Hui, B., Gao, B., Yu, B., Li, C., Liu, D.,
Tu, J., Zhou, J., Lin, J., Lu, K., Xue, M., Lin, R., Liu, T., Ren, X., and Zhang, Z. Qwen2.5-man m-athh technical report: Toward mathematical expert model via seff-improvement.
arXiv preprint X4X1:2449.11/222., 2024.
Ye, T., Xu, Z., Li, Y., and Allen-Zhu, Z. Physics of language models: Part 2.2, how to learn from mistakes on grade- school math problems. arXiv preprint arXiv:2408. 16293, 2024.
Ye, W., Zhang, Y., Abbeel, P., and Gao, Y. Become pro- ficient player with limited data through watching pure videos. In International Conference on Learning Repre- sentations, ICLR, 2023.
Yuan, S., Zhao, H., Du, Z., Ding, Liu, X., Cen, Y., Zou, X., Yang, Z., and Tang, J. Wudaocorpora: super large- scale chinese corpora for pre- pre-training language models. A/ Open, 2:65 68, 2021.
12

---

Controlling Large Language Model with Latent Action
A. Architecture of CoLA
A.1. Language World Model
The language world model fworld is the core component of language generation, aiming to predict the next token 2++1 from the current token sequence Z1:t under the latent actions dz. The design for the language world model is as follows:
Base Model: A large language model trained by standard auto-regression, which maps the token sequence Fia to embedding e/. Note that e/ also serves as the input embedding of the inverse dynamics model and policy model.
Merge Module: A simple module consisting of Nm. specialized MLPs, which we call merge-MLP and are similar to the intermediate layers of LLMs but modified to take as imput the concatenation of the embedding and the latent action: [e]. ar]. and output a new embedding et" of the same dimensionality as the original embedding c/. From the following merge MLP, we continue concatenate et and dr to input into the next one. Then an Im-head maps the final embedding of the next token distribution.
By the design of the language world model, we can transform an auto-regressive model into an action -itionn-ovvreedd world model with only a few additional parameters.
A.2. Policy Model
The policy model ı to output the latent action to guide the token distribution generated by the world model, which is the core component for RL. It is designed standard N, transformer blocks, but the output head has a size equal to the number of latent actions, which is the logits of each latent action. It takes the token sequence embeddings ef.e- which is from the base model in the language world model, input and outputs the distribution of the next action.
A.3. Inverse Dynamics Model
The inverse dynamics model fluverse aims to construct such latent action space for language models. With the world model and policy, we can build language model governed by latent action. However, we still face the challenge of determining how to obtain such a latent action space. Since we only have token- language data and no actual actions, we first need to consider and define how to extract the latent action. We think that latent actions can be inferred from the generated results. Thus, our design for extracting latent actions is an inverse dynamics style, which takes current state and future state as input to output the executed action (Tian et al., 2024). For the latent action space design, we employ discrete latent actions because prior research has shown that continuous latent action spaces suffer from problem known as "shortcuts" (Ye et al., 2023). In this issue, latent actions only capture information corresponding to the immediate next step, ignoring broader contextual information and hindering the ability of latent actions generalize well. Specifically, we adopt a codebook C { of size N, where each c, corresponds to specific candidate action. Thus, in our language framework, to predict the action as at time I, the inverse dynamics model takes as input the historical context . and future context #e+ lstte: to output the action dN. It contains two parts:
Encode Module. The encode module is constructed by N₁ blocks of causal transformer, which take the embedding of context and future Fitt+e (In our paper, we set to be 1) input, then takes the final position of mapped embedding ēt+o as output to serve the current time embedding et-
Action Mapping Module. Then the action mapping module maps the embedding ej to select a certain a, from C. For the latent action selection, traditional methods such as VQVAE (van den Oord et al., 2017) often adopt distance- based projection to match the action between embedding e; and codebook C, and then employ reparameterization tricks to ensure gradient propagation. However, in our experiments, we observed that this suffers from codebook collapse, where only a limited number of actions are activated during training. To address this issue, we redesigned the codebook projection mechanism. We implemented a direct code assignment approach. First, an action head maps the embeddings ei logits vector I of length N (the size of the codebook). Then, using Gumbel-Softmax sampling, we obtain one-hot vector of = OneHot where of = CumbelSoft (1t) and OneHot(-) means setting the largest value in the vector to and the remaining values to O. Note that since we use Humbel-Softmaa for sampling here, which samples a one-hot vector based on a softmax probability distribution, there is a certain probability of deviating from the optimal action assignment and selecting other actions. Then to guarantee the gradient backpropagation, we adopt a reparameterization trick to obtain a differentiable one-hot vector of:
13

---

Controlling Large Language Model with Latent Action
o, = (o) - g))ag + gi
(2)
where (:)ag means stop gradient. Finally, we construct a linear mapping from the codebook C: W. = [cf NN which projects the one-hot vector of into the codebook via matrix multiplication at Wo Finally, we map the current time embedding e, a latent action a E C.
A.4. Conclusion of Model Design
Finbaking Space
sample
sample
d1
a1 ar
Inverse Dynamies
Policy
ef....er..
e
e .1....e e
(b)
Language World Model
a1 .r
X1 , .I
Base Model
Merge Module
X2 XT+1
e
(c)
MLP
Transformer Block
Concatenate
Figure 6. The Model Structure of CoLA. (a) Inverse Dynamics Model: taking future conditioned context as input and outputting the latent action. (b) Policy Model: taking the context as input and outputting the latent action. Language World Model: taking the context and selected latent action as input to predict the next token.
We introduce the whole model structure in Figure All the embedding dimensions and other hlper-paramette in the transformer are the same those in Llama-3.1-8B as well the embedding dimension of code in the codebook. The merge module consists of multiple merge MLPs. Merge MIP is an MLP block similar to the intermediate layer in Llama-3.1 1-B. We introduce the forward process of that block: For the input embedding ef and selected action a,, we concatenate them by [e), ax] as input. First, two linears W, and W2 project the input to embeddings el and e? with the size of the intermediate size in Llama-3.1 1-B. Then compute the embedding e, (e!) @ ef. Finally, linear maps the embedding e, to éj with the same dimension as et. And éf serves as the input embedding of the following merge MI.P. Finally, we map the output embedding of the merge MLP the token logits.
A.5. Discussion Model Design
First, in terms structural design, the inverse dynamics model uses future conditioned information to extract latent control conditions. This allows us distinguish and identify distinct future generations based on different control conditions, which reduces the uncertainty of prediction. For the design of the language world model, since CoLA aims to separate high-level control from low-level language capabilities, pre-trained auto-regressive model, which inherently possesses basic language abilities, is well-suited to "actionize" the base language capabilities by inserting control conditions. This insertion is similar to multimodal models (Hong et al., 2023), From this perspective, we can view latent actions as a high-level modality we construct, which compresses and encapsulates abstract future information.
For the tuning of action and token modalities, we delegate tasks requiring basic language capability adjustments, such as
14

---

Controlling Large Language Model with Latent Action
instraction-followingg to the low-level module. Meanwhile, tasks involving alignment with high-level objectives, such as specific human preferences or intents, are handled by the high-level module. During training. these two components also assist each other. For instance, in SFT with action guidance, the high-level conditions reduce future uncertainty, enabling the low-level module to adjust more efficiently. During the RLHF phase, the fixed low-level language capabilities ensure stable and non-degrading language generation, making high-level learning more robust.
B. Training of CoLA
B.1. Model Training Process
After completing the model design, we further introduce how train these parts, including the inverse dynamics model parameterized by O the language world model parameterized by wworld (fbases meerge)) where 0base is the parame- terization of base model and meerge that of merge module, and the policy model parameterized by P We use 0 to denote frozen parameter. We divide the training into three stages: Constructing latent action control, tuning the language world model under action guidance, and latent action level reinforcement learning.
B.I.I. CONSTRUCTING LATENT ACTION CONTROL
In this stage, we introduce large corpus of dataset Dpre = {æır] } like pre-training to train all the newly added parameters. First, we jointly train the inverse dynamics model and language world model by:
min
Cpre1
min Pluverse e Omerge
Spredict + 8 &
Plaverse Amerrs
The first term Spredict to predict the next token:
Ei=n log Pworla (#e+1 Fias de merger basse)] }]] where de computed by f (#1:1+1. (lverrs)) with differentiable trick Equation 2, to guarantee the gradient backpropagation. Note that only the merge module in the language model is optimized. And the regularization term Lreg is an entropy regularization in the selection of codebook: Ee E Em= gi, log gi where gi is the k-th value of vector g. This regularization term can mitigate the issue of codebook collapse.
EuI. pre
Then we initialize the policy model to mimic the action selection of inverse dynamics by minimizing the objective:
4pee2 = - -ELITT
log x (Ce 11. pplllcy)
(4)
where a, = (2 Onverss).
1.1.2. FINE-TUNING UNDER ACTION GUIDANCE
In this stage, we utilize dataset Deft { Ip+1:T] formatted for instraction-following tasks, where Tiip is the instruction and Ep+1:7 is the response. We need to acquire cntrucctionfollovviin capabilities through such data. Since such formatted data did not appear during the last phase, we still need to fine-tune the language world model to adapt to the cnttruction-oollwviii guidanee (FTA), where we fix the inserted latent action and tune the base model to fit the instruction following mode. According to the source of latent actions, we propose two types of FTA in distinct scenarios.
mode. We propose a special method to tune our world model called Fine Tuning under Action
FTA from inverse model (FTA-I). When the dataset is diverse and contains multiple types of distribution, utilize FTA-I Similar to language world model learning in the pre-training stage, but without the regularization term, we optimize the world model by:
min min EgzyyroDaar
Abase
Buase
(5)
Σ log Pworld (Ie+ par at, Obaser mergg))]
15

---

Sample batch of æ1rr from Dat.- Controlling Large Language Model with Latent Action
Algorithm 1 Pre-Training
Algorithm 2 SFT
Input: Pretraining data Pre and iters Kree which com-
Input: SFT-TYPE (FTA-I, FTA-P), SFT data Dan and iterations Kytt
puted by the total training tokens.
# Step
for 1,..., Ksft do
fo 1,..,, K pre do
Sample batch of 21:7 from Datt-
Sample a batch of 21.7 from Prre-
Learn worll by Equation 5 with Dan. and SFT-TYPE. end for if SFT-TYPE FTA-I then
Learn world and iinverse by Equation 3. end for
Step 2 for t 1,..., pre do
for = 1, ...... K,n do
Sample batch of 21.r from D pre-
Compute action target a:77 by finverse. Learn 6 policy by Equation 4.
Compute action target a1.7 by finverse.
Learn policy by Equation 4.
end for end if
end for
Algorithm Roll-out
Algorithm 4 Reinforcement Learning
1: Input: Prompt (II, Tpp)
1: Input: Prompt (fı,..., Tp) and initial model "poliey
2: for t p...., .....,To do
2: Generate sentence F1.7 by Algorithm 3.
3: Select action a, by the cognitive policy;
3: Compute the reward by r(x1r)). and kl by initial model.
4: Sample the next token T++1 by the world model;
4: Optimize the policy model (opplsey to maximize r(Z1er) by an iteration of an RL algorithm.
5: end for
6: Return T1:T
5: Return policy model Topaisey
where de is computed by finverse (X1a+1, 0inverse) with freezed inverse dynamics model. We also freeze the merge block and only tune the base model to switch it to instruction following mode. After tuning the world model, since the embeddings provided by the base model have changed, we fine-tune the policy model by imitating the output of the inverse dynamics model, which is similar to Equation 4 but only imitating action corresponding to responses.
FTA with policy model (FTA--P). When the data distribution is narrow, such as mathematical reasoning datasets, we observe FTA-I will lead much lower loss in Objective 5, indicating overfitting to overly specific future outputs, akin shortcuts. Thus, we adopt FTA-P, which fine-tunes the language world model using Equation 5, but with the action provided by the policy model: = Spolicy (Zi:: Pp Since the action is provided by policy model, we do not need to further tune the policy to fit the changed embedding.
B. 1.3. LATENT ACTION REINFORCEMENT LEARNING
We further align the language generation with human preferences or other control goals by RL. In the RL stage, a prompt- only dataset D.L {.1ş} provided for sampling responses, and a reward model R(x1:)) which represents a specific preference goal, for reward signals. We optimize the policy model by maximizing the cumulative rewards:
max Exi, Dei.nraa 1r policy Swootd [R(xTT)) Opoliey
(6)
where we sample latent actions from the policy æ to input into the world model and select the token with the maximum probability from the world model's prediction.
B.2. Training Algorithm
We summarize the training algorithm of pre-training in Algorithm 1, the post-training in Algorithm 2, and the RLHF process in Algorithm 4. For RLHF algorithm, the KL divergence is computed on the latent action space, where the reference model the initial model of policy.
16

---

Controlling Large Language Model with Latent Action
C. MCTS Algorithm
The CoLA model, due to the smaller latent action space, reduces the search space, enabling more flexible control. Here, we present a latent action-level MCTS (Swiechowski al., 2023) approach called MCTS-Q for more efficient search. Compared with MCTS, MCTS-Q modifies the expansion steps by introducing Q-based model to provide value Q[ az) each time step t for pruning, ensuring that expanded search is only performed where necessary. We adopt Double-DON (van Hasselt et al., 2016) to learn the Q-function. Since our action space is much smaller than the token level, such Q- 2-function is casier to learn. After learning the Q-function, we define the uncertainty of a certain transition (I1:a, Qz, #1:c++) 1) by computing the Bellman error (Sutton & Barto, 1998). If the error is larger than a threshold, the current transition is defined as having large uncertainty: otherwise, it is defined as having low uncertainty. In MCTS-Q Q, when an expanded node, where the state is F1+++ and the token action from its parent is a,, computed with low uncertainty, we do not start the simulation. Instead, we continue take & step actions to generate and concatenate the generated tokens to the state until the node has large uncertainty.
The standard algorithm contains four steps:
Selection:
Start at the root node of the tree. Then traverse the tree by selecting the most promising child nodes based on a selection policy, such as UCT:
( In N()) N(v))
UCT(E;,
+ +oV ((v))
where U; is the child node being evaluated, v is the parent node of 0s (vx.) is the total reward accumulated from simulations passing through node v; N (ux) is the number of times node v; has been visited, N Nu) is the number of times the parent node v has been visited, c is a constant exploration parameter.
Expansion: When a leaf node is reached (a node that has not been fully explored), expand the tree by adding one or more child nodes on the selected leaf node. These child nodes represent possible actions from the current state.
Simulation: From the newly expanded node, perform a random simulation (rollout) until terminal state is reached. The result of the simulation is used to estimate the value of the node.
BackPropagation Update the statistics of all nodes along the path from the expanded node back to the root node, Increment the visit count (v) for each node. Update the total reward Q(v) based on the result of the simulation.
In our language generation, each node v contains the state, which is the historical context, the child set, which labeled by action reach the child, the value Q(v), and the visit count N(v). For each expanded node, we save its simulation content and final simulation value, which is obtained from the Qwen-2.5-Math-775 reward model. The action is multi-token sequence with fixed steps k. The state of the root node is the prompt. We repeat the MCTS step (from selection to BackPropagation as one step) for Nme- But the expanded node reaches the terminal (end token of the sentence), we can stop the MCTS early. After finishing the MCTS, we check all the nodes and select the nodes where their simulation value is the largest. We concatenate the state and its simulation content as the selected response.
Since our CoLA model has constructed the latent action space, we aim to apply MCTS on the latent action space, which may be more flexible. However, the latent action still only controls one-step token, which needs large cost of time. To save the time but search with flexibility, we introduce MCTS-Q, which introduces a learned Q function for uncertainty estimation and search pruning.
MCTS-Q algorithm. Compared with MCTS, MCTS-Q modifies the expansion steps by introducing a Q-based pruning. First, we introduce the learning of the Q function. The Q function is a Llama-3.1-88 model but replaces the Im-head with a linear layer from vocabulary size action size. Given a prompt set {#1p) from the math training dataset, we utilize the
CoLAA model after FTA-P generate N, responses (x++1:T) by sampling action sequence dp:r-1 for each prompt and model. Utilizing the dataset { Ip ap:T-1, r), we adopt Double-DQN DQN to learn the Q-function Qo parameterized by 6:
label the responses with reward {r} by the Qwen-2.5-Math-77
Lrmg (8) = T p t=p
(Y, a())) 0))
where the target Y, is computed by:
17

---

Controlling Large Language Model with Latent Action
Y = f Q arg max d);;; otherwise.
if 1ir terminal,
The target network Qo is updated by 0- 7 + (1 101--74. r)8- for every Ng gradient steps.
After learning the Q-function, define the uncertainty of a certain transition (#1=+11, az) by computing the bellman error (Y₁- - (z1: ax; 0))² where after training, 0 equals to 0. If the error is larger than threshold b, the current transition defined as having large uncertainty; otherwise, it is low uncertainty. In MCTS-Q, when an expanded node, where the state is T1:++1 and the token action from its parent is a, (or last step action), is computed with low uncertainty, we do not start the simulation. Instead, we continue to take k step actions to generate and concatenate the generated tokens to the state until the node has large uncertainty.
C.1. Training Details
For model design, we use 1lama-3.1-8P as the base model, additional N; = 4 transformer layers as the inverse dynamics model, Nm = 2 me as the merge module, and Np = 8 transformer layers as the policy model. For the number of codes, we use = 64 latent actions, where each code has the same dimension as the token embeddings.
C.I.I. DETAILS OF PRE-TRE TRAINING
We provide the details, including the hyperparameters and resources, during pre-training. For the pre-training hyper- parameters, adopt a learning rate of le length of 2048, and a maximum gradient norm of 1.0 for both the inverse dynamics model, the language world, and policy pre-training. For inverse dynamics model and language world model training, we adopt regularization loss, and its coefficient 3 is set to be 0.001. For the pre-training hyper-parameters in the ablation study, we set the learning rate to be le - 5 in the ablation of the dataset since we need to train all the parameters. For the ablation of parameters, since introduces the same trainable parameters, we keep the same hyperparameters. For evaluation, we utilize N& = 100 sequences with length of 2048 to compute the prediction loss, semantic diversity and KL computation. When computing generation semantic diversity, we need to take the prefix of the sequence for generation, the length of prefix is set be 256.
global batch size of 512, micro batch size of 4, maximum sequence
C.1.2. DETAILS OF POSTIIRAINING
We provide the details, including the hyperparameters and resources, during post-training. For the post-training hyper- parameters, first for SFT and FTA-I in preference tasks, we utilize learning rate with 5e - 6, training epoch with 1, global batch size with 256, and micro batch size with 4. Since we tune the same parameters as Llama-3.1-8B at this stage, the baseline adopts the same parameters as our CoLA model. For reward learning, we utilize BT model training based on Llama-3.1-88 model, learning rate with 9e - 6, training epoch with 4, global batch size with 256, and micro batch size with 4. The loss is computed by LogSiigmidd(((r y r(z, y ))) where x is the prompt, y is the chosen response and y is the rejected response. The reward model is utilized for both CoLA and baseline. For RLHF, we recommend using LLM--peccficc reinforcement learning methods, such as ReMax (Li et al., 2024b), RLOO (Ahmadian et al., 2024), GRPO (Shao et al., 2024), and REINFORCE+ (Hu, 2024), which all save memory and accelerate convergence. We use max generation length with 1024. For Math RL, the max length is set to be 2048. For agentic RI, we chose the validation task set of each environment to perform RL since the training set is too large. The max length is 4096. For FTA-P, we adopt the same hyperparameters as FTA-I, which the same as the baseline.
C.1.3. DETAILS OF MCTS AND MCTS-Q
We provide the details of MCTS and MCTS-Q. The algorithms provided in Appendix C. For MCTS, the length of multi-token search & is 64, the max repeating number Nmc is 64, and the coefficient in UCT is 0.7. For MCTS-Q, the threshold is set to be 0.01. For Q function learning, the learning rate is 5e - 6, the learning epoch is 100, the number of generated responses N, is 8, the update interval for target Q is 100, T is 1.0, the global batch size is 256 and the micro batch size is For the Q function learned in baseline, we only replace the output head with the vocabulary size but keep all the training hyper-parameters the same.
18

---

Controlling Large Language Model with Latent Action
Table Performance of CoLA and baseline on Benchmarks. The Base Model the initial model, and the FT model tuned certain domain dataset. ACA is academy, BUS is business, ENT is entertainment and LIT is literature. We mark the improvements of FT relative BASE red and declines blue. P-shift the parameter difference from tuned model the initial.
LLAMA-3.1
1-B8BB (BASE)
CoLA (OURS)
BENCHMARK
MMLU
GSM8K
MATHQА
P-SHIFT
MMLU
GSMSK
MATHQA
P-SHIFT
BASE MODEL
65.14
49.51
39.73
64.96
48.75
34.20
FT MODEL- MOA
64.85-0.29
40.56-9.95
38. 38.09-1.64
5.41
65. 65.12+0.16
52.01+3.26
34.87+0.67
4.72
FT MODEL-BIL BUS
65.08
0.06
28.28
21.23
38 89 -.844
0.84
5.39
65.13
49.20 +0.45
34.71 +0.51
4.76
FT MODEL-ENT
64.59
0.55
38.29
11.22
40.13
5.53
60.47
39.20
35.11
4.92
FT MODEL-LIT
0..5
37.60
-1.91
39.50
..23
5.69
65.07
50.49.1.74
36.18+1.00
5.57
0.38
13.58
-0.88
5.51
0.21
-1.03
1.02
4.99
D. Additional Empirical Results
D.1. Computational Overhead of Framework
We analyze the training parameters during Fine-Tuning and RLHF, including FTA-I and FTA-P, and the training parameters for RLHF. Results are shown in Figure (a), demonstrating that we introduce a small number of training parameters during the SFT stage, which nearly 1.25 times, but significantly fewer parameters during the RI., which is less than 0.25 times. Then we compare the time cost during tuning and RLHF, including the training time of the two tuning variants compared to standard tuning, the generation time in RLHF and the optimization time in RLHF, Results are shown in Figure 7 (b). demonstrates that only marginally increase the time for training and inference due to the additional parameters.
1.2
CoLA
baseline
1.2
baseline
1.0
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
FTA-1
FTA-P
RLHF
FTA-1
FTA-P
RLHF-gen
RLHF-opt
(a) Training Parameters
(b) Time Cost
Figure 7. Relative Cost of CoI.A model comparing with baseline Model. (a) is the relative numher of training parameters where baseline is to 1. (b) is the relative cost of time where RL.HF-gen means the generation time during RL.F and RLHF-opt means the optimization time cost during RLHF
D.2. Ablation Study on Dataset and Parameters
We aim ablate that the additional parameters and dataset can not attribute to the performance of the baseline model. For the ablation of dataset, we continue to train Llama-3.1-8B model for 1, 2 and 5G tokens the dataset. For the ablation of additional parameters, we add 8 transformer blocks to the end of Llama-3. 1.1ma-3311-0 1-8B transformer layers and freeze other parameters, which are the same trainable and inference parameters as our policy model in CoLA but only serve as the forward layer in the nugo-regressive model. We evaluate the semantic diversity and MMLU value. The results Figure show that with the increasing training tokens and parameters on the dataset, the base model shows a continuous decrease in performance, indicating that the additional dataset and parameters can not be attributed the performance of the baseline.
19

---

Controlling Large Language Model with Latent Action
1.34
Larra-3.1.88 Ablation on Params
Liama-3.1-08
1.53
Ablation
Ablation on Params
64.8
1.52
Value
Value
64.6
1.50
64.2
1.49
64.0
63.8
1.0
1.5
2.0
2.5
3.0
3.5
4.0
4.5
5.0
1.0
1.5
2.0
2.5
3.0
3.5
4.0
4.5
5.0
Num Training Tokens
Num Training Tokens
(a) Performance Diversity
(b) Performance MMLU
Figure 8. Performance of ablation on dataset and parameters. Green line is the ablation on parameters, training from IB to 5B tokens. The yellow line is the dataset, training from 1G 5G tokens. Blue line is the baseline 1lama-3.1 1-8B model.
D.3. Ablation on Components in CoLA
CodeBook Learning Method. We compare our direct action assignment to the traditional VQVAE methods. We counted the number of times each of the 64 actions was used during training and calculated the number of actions used more than times, which refer alive actions. Figure 9 shows that the VQVAE suffers great code collapse the alive actions are much lower while our direct action assignment can be stable.
Ablation between FTA-I and FTA-P. We also compare the math500 performance of FTA-P with FTA-I and SFT of baseline. After training on the same dataset NuminaMath, the greedy performance of baseline on math500 36.0, where CoLA model with FTA-I is only 25.0 and FTA-P is 41.0. indicating that FTA-P can be more suitable on such tuning dataset.
Selection of Distinet Base Model. To demonstrate the scalability of our approach across different base models, also tested the effectiveness of the CoLA design on the Qwen-2.5-Math-115. model. We adopted the same layer design and introduced additional 0.5B parameters to the 1.5B model. We observed in Figure 10 that the loss still effectively decreases during the pre-training phase, and the codebook does not collapse, indicating effective latent action control learning.
D.4. Latent Action Control
Quality under Latent Action Control. We compute the generation quality by caleulating the quality value, which is directly evaluated using the Qurator model (Wettig et al., 2024). The results in Figure 11 demonstrate that our latent action space achieves better generation quality compared to the token space.
We evaluate on the validation dataset Dval of pre-training phase. We compute the prediction loss Spredict in Equation on D The loss 0.45, which is much lower than that of Llama-3. 1-8B, which 1.77. This indicates that with accurate latent actions, the predictive uncertainty can be reduced through latent action.
Uncertainty under Latent Action Control.
Potential Relationship between World Model and Auto-Regresse Model. After pre-training, if we consider these action-controlled distributions as marginal distributions of the raw next token distribution of base model, we calculate the expected distribution of the latent actions:
N
(a = c2 p(1+11221)) Σ
c;)
where E C. Comparing with the corresponding next token distribution by the base model, the KL distance is 0.09, indicating that our latent actions potentially decompose the original token-only distribution.
Visualization Latent Action Control. Then we visualize the words generated under different latent actions using word
20

---

Controlling Large Language Model with Latent Action
60
Num Alive Action
50
40
30
20
10
Direct action assignment
Vq-Vae action assignment
0
0
100
200
300
400
500
600
Num Gradient Steps
Figure 9. The alive action number during training. The blue line is the direct action assignment, and the yellow line vqvae.
clouds (Kalmukov, 2021). Results of several latent actions shown in Figure 12. This reveals that the latent actions exhibit a degree of interpretabiliity. For example, the action 1 controls significant number of rrogramming-relatee terms such "java", "python", "ruby" and others. And the action 2 controls some time-related terms and some names of people and places. In most cases, it simply groups words with similar meanings together. For example, action 3 includes words like "long" and "major", while action 4 includes "give", "offering" and "giving", or "guide" and "guidance" among others. This has indicated that our latent actions have learned some higher-level meanings beyond individual tokens. The actions can effectively govern distinet token distributions and show degree of interpretability. These provide foundations for further controllable tuning and alignment in post- post-training training stage.
D.5. Generation Results in Distinet Tasks
Generations on RLHF with kl=0.00. We show more results generated by the model after RLHF with KI. coefficient 0.00. For distinct preference, baseline always falls into fixed pattern, and under certain preferences, issues such generating irrelevant content repetitive language degeneration arise, indicating reward hacking problem:
Academy Preference:
21

---

Controlling Large Language Model with Latent Action
16
Num Alive Action
50
Prediction Loss
20
10
50
75
100
125
150
175
25
50
75
100
125
150
175
Num Gradient Steps
Num Gradient Steps
(a) Num of Alive Action
(b) Training Loss
Figure 10. Pre- re-Training on Qwen-2.5-Math- Model. (a) is the number of alive action. (b) is the training loss.
Instruction Find three examples of British English slang.
CoLA: 1. "Bollocks" Meaning nonsense nonsense talk
2. "Chuffed" Meaning very pleased or satisfied
3. "Knocking on" Meaning trying hard or trying to succeed at something Baseline: As a researcher, would like to clarify what you mean by "British English slang. Could you please provide some examples?
Instruction 2: For new employee, what duties should be included in their
job description?
The job description for a new employee should clearly define and outline the responsibilities, duties, and functions that are required for the position. This should include the main purpose the role, key performance indicators, and list of specific tasks and activities the employee is expected to perform. It should be comprehensive enough to give them sense of what their day-to-day work will entail and should be relevant to the job title and organization's objectives. In addition, it is essential to include details on the required skills, experience, education, and
CoLA:
for the position. The job description should be concise, clear, and accurate, and should be reviewed and updated regularly. Baseline: As researcher, would like to clarify what you mean by "duties" that should be included in new employee's job description. Could you please provide an example?
22

---

Controlling Large Language Model with Latent Action
CoLA baseline
-2
-4
-5
Random Sampling
Model Sampling
Figure 11. The quality value. The left part is the random sampling on the token or action space (CoLA means action space and baseline means token space). right part the model sampling, where CoLA means sampling by policy model and baseline means sampling by 1Jama-3.1-8B.
Handon yubyy
specifu
fer
Lude
fic
ent
Donning Paul
Smat
subsequeng
Leweer
ong
Sealler
moough
precent
According
Corsing
samp le
spaces
source CD
8 enought
Vving
splece ? off
C cLass
softwar
Java
Tuesda
ma
Scutheavy
nigher
Wakinngg
Javascr wwt.crrptt.......
Python
scala
Laaioo tersoneel
televiside sn
ooday
Leng h
mammmshhh ish
usual
dancee Sorrettti
(a) Latent Action
(b) Latent Action 2
(c) Latent Action 3
(d) Latent Action 4
Figure The word cloud the words controlled by distinct latent actions
Instruction 3:
Generate list of tips for managing stress during exams.
CoLA:
"Sure! Here are some tips for managing stress during exams:
1.
Create schedule: Make a schedule for each day leading up to your exams to keep track of your time.
2.
Take breaks: Make sure you're re taking breaks in between studying to keep
your brain refreshed and focused.
3. Exercise: Physical activity can help reduce stress and improve mental
health.
4. Eat healthy diet: Eating a balanced diet can help you stay alert and
focused.
5. Stay hydrated: Drinking water helps keep your mind and body alert.
6. Stay organized: Stay organized during your studies to prevent stress
from piling up.
7. Take deep breaths: Deep breathing exercises can help reduce stress and
anxiety.
8. Talk to someone: Talking to family, friends, professional help
you manage stress.
9.
Get enough sleep: Getting enough sleep can help you stay focused during
your studies.
10. Practice relaxation techniques: Practicing relaxation techniques like meditation and yoga can help you manage stress and improve your overall health.
Baseline: As researcher, would like clarify what you mean by
"managing stress during exams." " Could you please provide some examples?
23

---

Controlling Large Language Model with Latent Action
Instruction 4: Translate the following sentence from French to English: "Je veux lire livre."
COLA: "I want to read book."
Baseline: As researcher, would like clarify what you mean by "read book." " Could you please provide an example?
Business Preference:
Instruction 1: What is benefit of sending automated emails?
Sending automated emails have many benefits for businesses. Here are some of the most common advantages:
CoLA:
1. Increased efficiency: Automated emails allow businesses to send messages at scale without the need for large team of employees. This can save a lot of time and resources, allowing businesses to focus on more critical tasks.
2. Cost savings: Sending manual emails can be time-consuming, requiring a large team of employees. This can be expensive for businesses, particularly if they need to send large number of emails. By using automated emails, businesses can save on costs while achieving the same results.
3. Personalized messaging: Automated emails allow businesses to personalize their messaging by including recipients' names, location, or other relevant details. This can increase the effectiveness and response rate of the messages.
4. Improved engagement: Automated emails help businesses stay in touch products, services, and
with their audience and keep them informed about promotions. This can increase engagement rates and drive sales.
5. Improved customer service: Automated emails can be used to provide support and assistance to customers. For example, automated email responses can be used to answer common customer queries, reducing the workload on the customer service team.
Overall, automated emails are a powerful tool for businesses looking improve their efficiency, save costs,
Baseline: As professicnal corporate manager, I understand that sending automated emails can provide several benefits. Here are some of the key advantages of sending automated emails:
1. Increased Efficiency: Automated emails can help save time and effort by eliminating the need to manually create and send emails. This can help improve efficiency and productivity, allowing your team focus on other critical tasks and projects.
2. Improved Accuracy: Automated emails help ensure that all emails are sent with precision and accuracy, reducing the risk of errors and mistakes. This can help improve the quality of your work and enhance customer satisfaction.
3. Better Engagement: Automated emails can help engage customers and stakeholders by providing them with relevant and timely information. This can help improve their experience and drive success. As corporate manager, it's essential to use automated emails to achieve success, efficiency, and productivity.
24

---

Controlling Large Language Model with Latent Action
Instruction 2: What is the best way to engage with customers social media?
The best way to engage with customers on sccial media is to be zesponsive, empathetic, and helpful. It's important to be proactive in engaging with your customer base, rather than waiting for them to reach out to you.
COLA:
You should also make promptly to any questions, concerns, complaints that they may have. Showing that you care, and are willing to go above and beyond to provide excellent customer service, can help you to build strong relationships with your customers.
to listen carefully to your audience, and respond
Additionally, it's important to be creative, and think outside the box when engaging with your customers on social media. Offering exclusive deals and promotions, hosting fun and interesting contests can help you to stand out from the crowd, and build more engaged audience.
At the end of the day, the best way to engage with customers on social media will depend on your business, your brand, and your audience. It's important to stay up-to-date with trends and best practices, and be willing to try new things to find what worked best for you. Baseline:
professional corporate manager, the best way to engage with customers social media be achieved by follewing these steps: 1. Establish strong online presence: It's essential to have visible and engaging online presence social media platforms. This can help you build your brand's reputation, reach new customers, and improve your overall success.
2. Listen to your customers: It's critical to listen to your customers' feedback and concerns on social media. This can help you improve your products and services, resolve issues, and build stronger relationships with your audience.
3. Provide exceptional customer service: Engaging with customers social media can help you provide exceptional customer service. This can include addressing their concerns, providing solutions, and improving their experience. It's essential to create a culture that prioritizes customer satisfaction and success.
25

---

Controlling Large Language Model with Latent Action
ormulate research question related to the environmental impact of increasing population density.
Instruction 3:
CoLA: How does increasing population density affect the environmental
impact of urban areas in terms of energy consumption and carbon emissions?
Baseline: As a professions corporate manager, it's essential to understand the environmental impact increasing population density. To formulate research question related to the environmental impact of increasing population density, the following steps can be taken:
1. Define the problem: It's essential to understand the problem before formulating research question. As the population density increases, it can lead to various environmental impacts, such as pollution, zesource depletion, and climate change. It's crucial to address these impacts and take steps mitigate their impact.
2. Develop research question: As manager, it's essential to develop research question that can help address the environmental impact of increasing population density. The research question should be focused on understanding the impact and taking steps to mitigate its impact. It's essential to work with staksholders, experts, and the community to create a sustainable and@nate.
26

---

Controlling Large Language Model with Latent Action
Instruction 4: Design algorithm to solve interesting problem
I'm sorry, but as an AI language model I do not have the ability to design complex algorithms or solve interesting problems. However, I can provide you with the steps to create an algorithm to solve any problem: Define the problem: Tdentify the problem you are trying to solve and the outcome you want to achieve.
CoLA:
1.
2. Gather data: Collect the necessary information and data related to the problem.
3. Develop the model: Develop a model that represents the problen and the relationship between different variables.
4. Test the model: Test the algorithm by feeding different datasets to the
model and verify if it is providing the expected results.
5. Refine and iterate: Refine the algorithm based on the results and continue testing and iterating until the model provides optimal results. 6. Implement the algorithm: Once satisfied with the outcome of the model, implement it into a system or software to solve real-world problems. I hope this helps.
Baseline: As a professional onall corporate manager, it's s essential to design
an algorithm to solve an interesting problem. Here are the steps to follow when designing an algorithm to solve an interesting problem:
algorithm solve interesting problem is to define the problem. This includes understanding the problem's scope, requirements, and constraints. It's essential to have a clear understanding of the problem before moving forward with the algorithm.
1. Define the problem: The first step in designing
2. Analyze the data: As a corporate manager, it's crucial to analyze the data related to the interesting problem. This can help you understand the problem's complexity, challenges, and solutions. It's essential to use the data to make informed decisions and improve the algorithm's performance.
3. Test the algorithm: It's essential to test the algorithm before implementation. This can help you identify any issues and improve its performance. As a corporate manager, it's your responsibility to ensure that the algorithm is reliable, safe, and productive.
Entertainment Preference:
27

---

Controlling Large Language Model with Latent Action
Instruction 1: Explain the process of bacterial transformation
Ah, bacterial transformation! It's like giving new look to your old clothes. But in this case, we' re talking about bacteria. The process of bacterial transformation is a method used to introduce foreign DNA into a bacteria cell, typically for the purposes of genetic enginesring. First, you need to have bacterium that is sensitive the DNA you want to transform. This means that the bacteria's geneme will accept and incorporate the foreign DNA into its Next, you need to prepare the DNA that you want to transform. It's usually in the form of a plasmid, which is a small circular piece of DNA.
CoLA:
genome.
To transform your bacteria, you need to treat your bacteria cells with a chemical that temporarily makes their cell membrane permeable. Then, you mix the DNA plasmid with the bacteria cells and put them in a special device called electroporaatin chamber. The device uses small electric shock to open up the bacteria's s membrane, which allows the DNA plasmid to pass inside the bacteria.
Once inside, the plasmid is replicated and passed on to new cells when the bacteria divide. And voilà! Your bacteria now has foreign DNA. of course, there's a bit more to it than
Baseline: Well, that's an easy one! Bacterial transformation is like giving new genetic material to bacterium, and then watching it turn into totally different creature. Think of it like giving dog coat of fur and watching it turn into a cat. Just don't ask me to bring you bacterial transformation party, because that's not happening anytime soon!
Instruction 2: Create persuasive advertisement for running shoes.
Welcome to our show today! Are you ready to take your running experience to the next level? Do you need the ultimate confort and support for your feet while you run? Look no further! We have the perfect solution { our brand new range of running shoes that will take care of your feet and help you run like a champ!
CoLA:
Our running shoes are designed keeping all the needs and requirements of runners in mind. We understand the struggles you face while running, and our shoes
crafted to provide maximum support, comfort, and traction, ensuring that you stay comfortable, safe, and focused throughout your run. Our shoes are designed with advanced technology that offers a oushioned and responsive ride, helping you run for a long time and avoid injuries.
Our running shoes are made using high-quality materials that are durable and long-lasting. We understand that runners need shoes that they can rely on, and shoes won't disappoint you. Our shoes not just about comfort and support; they also look great, with stylish and sophisticated designs that will enhance your style quotient.
So what are you waiting for? Invest in our range of running shoes and enjoy the comfort, support, and performance you deserve while running! Don't let your feet hold you back; let them propel you forward! Get set Baseline: Well, that's an easy one! You just need pair of running shoes and beautiful view of mountain trail. Show up to the trailhead in your new shoes and get ready for the views of your life! Just don't forget the selfie stick on your way back down!
28

---

Controlling Large Language Model with Latent Action
Instruction 3:
Generate paragraph describing summer day in the park
On hot summer day in the park, the sun is shining bright in the sky, and the leaves of the trees are dancing in the wind. People of all ages can be seen lounging in the grass, basking in the warmth of the sun. Families can be heard laughing and chatting happily, while children run and play, creating a cheerful ambiance. The sweet smell freshly cut grass and flowers fills the air, complete with the occasional whiff of popeorn from street vendor. It's day of relaxation and enjoyment in the beauty of Mother Nature.
CoLA:
Baseline: Well, that's s an easy one! Imagine a beautiful day in the park, with the sun shining down and a cool breeze blowing through the trees. You hear the sound of children playing on the playground and see people laying out on blankets and reading books. You see dogs running around and people walking their pets. You smell the fresh cut grass and hear the sound of birds chirping in the background. Just don't ask to hold the hot chocolate, because there's none of that in sight!
Instruction 4: Make car analogy for the term "transaction"
CoLA: A transaction is like taking a ride in a car. You get into the car, the car takes you to your destination, and then the car drops you off and you get out. But, just like a car ride, the quality of the transaction depends on
driver. If the driver is skilled, efficient, and trustworthy, the ride is smooth and enjoyable. the driver is unreliable, unskilled, and unsatisfactory.
the other hand, if untrustworthy, the ride may be bumpy So, in the world of transactions, it's important to find reliable and trustworthy drivers who will take you to your destination safely and efficiently.
Baseline: Well, that's s an easy one! A transaction is like taking your car to the gas station and filling up the tank. You give the atterdant your money, and they put the gas in your car. Then you drive away and enjoy the next adventure. Just don't forget to pay the next time you stop the gas station!
Literature Preference:
Instruction 1: Give two possible consequences of global warming.
Oh warming earth, what wee you have brought, With rising seas, and storms that roar, The world's climate is a state of turmoil, With catastrophic effects, we must take note.
CoLA:
The first consequence that comes to mind, Is the displacement of people and their lives undone, As the seas rise, and the weather grows wilder, Their homes and lands are lost, and they have to migrate.
29

---

Controlling Large Language Model with Latent Action
second consequence, that we can't ignore, Is the extinction of creatures, large and small, As the climate changes, their homes become unviable, Leaving them no choice, but to breathe their last, and fall.
So let us heed this call, and take action, To save our planet, and all that reside, For it is the only home, have, to call our own, Let's work together, and make difference, and be proud. Baseline: Oh, the consequences of global warming, A menace that threatens all life and all things. With rising temperatures and changing patterns, The effects are rea-hig and quite alarming.
One consequence is the melting of glaciers, And the resulting flooding and displacement. As sea levels rise, coastal cities will be submerged, And the fate of the islands will be mourned.
The heat will affect agriculture and food supply, And the patterns of disease will shift and fly. With extreme weather events becoming more frequent, The consequences of global warming will be severe.
Oh, the consequences of global warming, A menace that threatens all life and all things. Oh, the consequences of global warming, A menace that threatens all life and all things. Oh, the consequences of global warming, A menace that threatens all life and all things. Oh, the consequences of global warming, A menace that threatens all life and all things. Oh, the consequences of global warming, A menace that threatens all life and all things. Oh, the consequences of global warming, A menace that threatens all life and all things. Oh, the consequences of global warming, A menace that threatens all life and all things.
Oh, the consequences of global warming, A menace that threatens all life and all things.
30

---

Controlling Large Language Model with Latent Action
Instruction 2: Name two items that are nearly impossible to find in stores during the pandemic.
Amidst the pandemic, when we are advised to stay indoors, And the world outside seems so chaotic and bitter, It's s hard to find some basic recessities, And the problem gets worse as time passes.
CoLA:
Two items that are hard to find, Are the ones that you need every day, Like toilet paper and hand sanitizer, The shelves are empty, there's s nothing to say.
These items, we used to buy without thinking twice, Are now luxury, hard come by, So l be kind to each other and community, And stay safe, healthy and positive. Baseline: Oh, the pandemic hath brought forth a great upheaval, In the world of commerce, sudden and abrupt cessation. With the shutters down, the shelves remain empty, Leaving us with a void that hard to ignore entirely.
In this time of crisis, there are certain items that scarce, And hard to find in the markets, leaving with sense of despair. One such item is the medical supplies that are in high demand, And the other is the food items that are essential for all and all.
The medi cal supplies like masks, gloves, and respirators, Are essential for the healthcare workers who are fighting the pandemicisors. And the food items like rice, wheat, sugar, oil, and pulses, Are essential for all vš vš vå va vš vš vš vá va va vž vš vš va vá Vă Vš Vš Vš vs vs vš Vš vš vs vs Vs Vá Vš Vššnate Wehaecceeeeeeeeeeeeeeeeeeeeeeeee Baacehaateenreehattcettttteeeee.
Generation on Agentic Tasks on Alfworld after RL.
Task Generation on Agentic Tasks on Alfworld after RL.
CoLA: user:
Interact with household solve task. Imagine you an intelligent agent in household environment and your target is to perform actions to complete the task goal. At the beginning of your interactions, you will be given the detailed description of the current environment and your goal to accomplish.
For each your turn, you will be given the observation the last turn. You should first think about the current condition and plan for your future actions, and then output your action in this turn. Your output must strictly follow this format:
31

---

Controlling Large Language Model with Latent Action
Thought: your thoughts.
Action: your next action.
The available actions are:
1. go {recep]
2. task {obj} from {recep]
3. put {obj} in/on {recep]
4. open [recep]
5. close {recep}
6. toggle {obi} frecep]
7. clean {obj} with {recep}
8. heat {obj} with frecep)
9. cool {obi} with {recep)
where {obj} and (recep) correspond to objects and receptacles.
After your each turn, the environment will give you immediate feedback based on which you plan your next few steps. if the environment output "Nothing happened", that means the previous action is invalid and you should try more options.
Your response should use the following format: Thought: <your thoughts> Action: <your next action>
assistant:
OK
user:
You are in the middle of a room. Looking quickly around you, you see a bathtubbasin 1, a countertop 1, a drawer 8, a drawer 7, a drawer 6, drawer 5, drawer 4, drawer 3, a drawer 2, a drawer 1, a garbagecan 1, handtowelholder 1, sinkbasin 1, toilet 1, toiletpaperhanger 1, towelholder 2, and a towelholder 1.
Your task is to: find two tissuebox and put them in toilet.
assistant:
Thought: need to locate two tissue boxes as per the task. The first place to check for a tissue box is drawer 1, as it is common storage place for such items.
Action: go drawer
user:
observation:
The drawer 1 is closed.
assistant:
Thought: The drawer is closed. I must open it to see if there's s tissue box inside.
Action: open drawer
user:
Observation: You open the drawer 1. The drawer 1 is open. In it, you see nothing.
assistant:
Thought: Drawer 1 is empty. I should close it and check the next drawer for a tissue box.
Action:
close drawer
user:
Observation: You close the drawer 1.
assistant:
Thought:
The next logical step is to check drawer 2 for a tissue box.
32

---

Controlling Large Language Model with Latent Action
Action:
go to drawer
user:
Observation: The drawer 2 is closed.
assistant:
Thought: This drawer is also closed. I'11 open it to search for a tissue box.
Action:
open drawer
user:
Observation: You open the drawer 2. The drawer 2 is open. In it, you see nothing.
assistant:
Thought: No tissue box in drawer 2 either. I should close this drawer and continue the search.
Action:
close drawer
user:
Observation: You close the drawer 2.
assistant:
Thought: Moving on to drawer 3 to find tissue box.
Action: go to drawer
user:
Observation: The drawer closed.
assistant:
Thought: Another closed drawer. I'll open it to look for tissue box.
Action: open drawer
user:
Observation: You open the drawer 3. The drawer 3 is open. In it, you see a candle 2.
assistant:
Thought: This drawer contains a candle, not a tissue box. should close
it and keep looking.
Action: close drawer
user:
Observation: You close the drawer 3.
assistant:
Thought: I'll check drawer 4 next for a tissue box.
Action: go to drawer
user:
Observation: On the drawer 4, you see toiletpaper 2.
assistant:
Thought: There's toilet paper on this drawer, but no tissue box. I'll move on to the next drawer.
Action: go to drawer
user:
Observation: The drawer is closed.
assistant:
Thought: Another closed drawer. I'll open it to search for tissue box.
Action: open drawer 5
user:
Observation: You open the drawer 5, The drawer 5 is open. In it, you see a tissuebox 4.
assistant:
Thought: found tissue box in drawer 5. I need to take it to complete the task.
33

---

Controlling Large Language Model with Latent Action
Action:
take tissuebox from drawer
user:
Observation: You pick up the tissuebox from the drawer 5.
assistant:
Thought: With one tissue box in hand, should close the drawer to keep the room tidy.
Action: close drawer
user:
observation: You close the drawer 5.
assistant:
Thought: I should place the tissue box in the toilet as instructed. I'11 go to the toilet to do this.
Action: go to toilet
user:
On the toilet 1, you see candle 1, and toiletpaper 1. assistant:
Observation:
Thought: I'm at the toilet. I'll place the tissue box here as required. put tissuebox in/on toilet 1
Action:
user:
Observation: You put the rissuebox 4 in/on the toilet 1.
assistant:
Thought: need to find one more tissue box. Drawer is the next logical place to check.
Action:
go to drawer
user:
Observation: The drawer 6 is closed.
assistant:
Thought: I'll open drawer to see if there's another tissue box inside.
Action: open drawer
user:
Observation: You open the drawer 6. The drawer 6 is open. In it, you see nothing.
assistant:
Thought: Drawer 6 is empty. I'll close it and continue the search.
Action: close drawer
user:
Observation: You close the drawer 6.
assistant:
Thought: Moving to drawer 7 to find the second tissue box. to drawer 7
Action:
user:
Observation: The drawer is closed.
assistant:
Thought: need to open drawer check for a tissue box.
Action: open drawer
user:
Observation: You open the drawer 7. The drawer 7 is open. In it, you see nothing.
assistant:
Thought: No tissue box drawer I'll close it and check the last drawer.
Action: close drawer
user:
34

---

Controlling Large Language Model with Latent Action
Observation:
You close the drawer 7
assistant:
Thought: The last drawer to check is drawer B. Hopefully, the second tissue box is there.
Action: go to drawer
user:
Observation:
The drawer is closed.
assistant:
Thought: I'll open drawer 8 to search for the second tissue box.
Actien: open drawer
user:
Observation: You open the drawer 8. The drawer 8 is open. In it, you see nothing.
assistant:
Thought: Drawer 8 is also empty. I'll close it and need to look elsewhere.
Action: close drawer
user:
Observation: You close the drawer 8.
assistant:
Thought: haven't checked the countertop yet. It's common place for
such items.
Action:
go countertop
user:
Observation: On the countertop 1, you soapbottle 1, spraybottle 2, spraybottle 1, a tissuebox 3, a tissuebox 2, and a tissuebox 1.
assistant:
Thought: There are tissue boxes on the countertop. I'll take one of them to complete the task.
Action: take tissuebox from countertop
user:
Observation: You pick up the tissuebox from the countertop 1.
assistant:
Thought: With the second tissue box in hand, I should close the drawer to keep the room tidy.
Action:
close drawer
user:
Observation: Nothing happens.
assistant:
Thought: have both tissue boxes now, need to place the second one in the toilet as instructed.
Action: go to toilet
user:
Observation: On the toilet 1, you see candle 1, tissuebox 4, and toiletpaper 1.
assistant:
Thought:
I'm at the toilet with the second tissue box. I'll place it as
required.
Action:
put tissuebox 3 in/on toilet 1
Baseline: user:
Interact with household to solve task. Imagine you are an intelligent agent in household environment and your target is to perform actions to
35

---

Controlling Large Language Model with Latent Action
complete the task goal. the beginning of your interactions, you will be given the detailed description of the current environment and your goal to accomplish.
For each of your turn, you will be given the observation of the last turn. You should first think about the current condition and plan for your future actions, and then output your action in this turn. Your output must strictly follow this format:
Thought: your thoughts.
Action: your next action.
The available actions are:
1. go {recep]
2. task {obj} from {recep)
3. put {obj} in/on frecep)
4. open [recep]
5. close {recep}
6. toggle {obj} frecep)
7. clean {obj] with {recep}
8. heat [obj] with {recep)
9. cool {obi} with {recep]
where {obj} and (recep) correspond to objects and receptacles.
After your each turn, the environment will give you immediate feedback based on which you plan your next few steps. if the environment output "Nothing happened", that means the previous action is invalid and you should try more options.
Your response should use the following format: Thought: <your thoughts> Action: <your next action>
assistant:
OK
user:
You are in the middle of room. Looking quickly around you, you see a bathtubbasin 1, a countertop 1, a drawer 8, a drawer 7, a drawer 6, a drawer 5, drawer 4, drawer 3, a drawer 2, a drawer 1, a garbagecan 1, handtowelholler 1, sinkbasin towelholder 2, and a towelholder 1.
toilet 1, tolletpaperhanger 1,
Your task is to: find two tissuebox and put them in toilet.
assistant:
Thought: need to locate two tissue boxes to complete the task. should start by checking the nearest storage areas, which are the drawers.
Action: go to drawer
user:
Observation: The drawer is closed.
assistant:
Thought: Drawer 1 is closed. I need to open it to see if there are any tissue boxes inside.
Action: open drawer
user:
Observation: You open the drawer The drawer is open. In it, you see nothing.
assistant:
36

---

Controlling Large Language Model with Latent Action
Thought: The drawer is empty. I should close it and check the next drawer for the tissue boxes.
Action: close drawer 1
user:
Observation: You close the drawer 1.
assistant:
Thought: I'11 move on the next drawer to continue my search for the
tissue boxes.
Action: go to drawer
user:
Observation: The drawer 2 is closed.
assistant:
Thought: Drawer 2 is also closed. need to open it to check for the tissue boxes.
Action:
open drawer
user:
Observation: You open the drawer 2. The drawer 2 is open. In it, you see nothing.
assistant:
Thought: This drawer is empty as well. I'll close it and continue searching the remaining drawers.
Action:
close drawer
user:
Observation: You close the drawer 2.
assistant:
Thought: I'll proceed to the next drawer to continue search for the
tissue boxes.
Action:
go to drawer
user:
Observation:
The drawer is closed.
assistant:
Thought: Drawer 3 is closed. I need to open it to look for the tissue
boxes.
Action: open drawer
user:
Observation: You open the drawer The drawer is open. In it, you see candle 2.
assistant:
Thought: There's only candle in this drawer. I'll close it and move to the next drawer.
Action: close drawer
user:
Observation: You close the drawer
assistant:
Thought: I'11 check the next drawer for the tissue boxes.
Action: go to drawer 4
user:
Observation: On the drawer 4, you see toiletpaper 2.
assistant:
Thought: There's s a roll toilet paper on drawer 4, but no tissue boxes. I'll close this drawer and continue searching.
Action: close drawer
user:
37

---

Controlling Large Language Model with Latent Action
Observation: Nothing happens.
assistant:
Thought: I'11 check the next drawer for the tissue boxes.
Action: to drawer
user:
Observation: The drawer is closed.
assistant:
Thought: Drawer 5 is closed. I need open it to if there any tissue boxes inside.
Action: open drawer 5
user:
Observation: You open the drawer 5. The drawer 5 is open. In it, you see a tissuebox 4.
assistant:
Thought: found a tissue box in drawer 5. I'll take it so can place it in the toilet as instructed.
Action:
take tissuebox from drawer
user:
Observation: You pick up the tissuebox from the drawer 5.
assistant:
Thought: With one tissue box in hand, need to place in the toilet. as per the task.
Action: go to toilet
user:
observation: On the toilet 1, you see candle 1, and toiletpaper 1. assistant:
Thought: I'm at the toilet. I should place the tissue box here to complete
part of the task.
Action: put tissuebox in/on toilet
user:
Observation: You put the tissuebox 4 in/on the toilet 1.
assistant:
Thought: need to find the second tissue box. I'll check the next drawer.
Action: go to drawer
user:
Observation: The drawer closed.
assistant:
Thought: Drawer 6 is closed. I need to open it to look for the second tissce box.
Action: open drawer 6
user:
Observation: You open the drawer 6. The drawer 6 is open. In it, you see nothing.
assistant:
Thought: This drawer is empty. I'll close it and continue searching the remaining drawers.
Action: close drawer 6
user:
Observation: You close the drawer 6.
assistant:
Thought: I'll check the next drawer for the second tissue box.
Action: go to drawer
user:
38

---

Controlling Large Language Model with Latent Action
Observation:
The drawer is closed.
assistant:
Thought: Drawer 7 is closed. I need to open it to check for the second tissue box.
Action: open drawer
user:
Observation: You open the drawer The drawer is open. In it, you see nothing.
assistant:
Thought: This drawer is also empty. I'll close it and continue searching the last drawer.
Action: close drawer
user:
Observation: You close the drawer
assistant:
Thought: I'l1 check the last drawer for the second tissue box.
Action: go to drawer 8
user:
Observation: The drawer is closed.
assistant:
Thought: Drawer 8 is closed. I need to open it to see if there's s the second tissue box inside.
Action: open drawer
user:
Observation: You open the drawer 8, The drawer 8 is open. In it, you see nothing.
assistant:
Thought: This drawer is empty as well. I1l close and recheck the previous drawers in case missed something.
Action: close drawer
user:
Observation: You close the drawer 8.
assistant:
Thought: I'11 go back to drawer 1 to double-check for the second tissue
box.
Action:
go to drawer
user:
Observation: The drawer is closed.
assistant:
Thought: Drawer 1 is closed again. I need to open it to check for the second tissue box.
Action: open drawer
user:
Observation: You open the drawer The drawer is open. In it, you see nothing.
assistant:
Thought: This drawer is still empty. I'll close it and continue searching the other drawers.
Action:
close drawer
Generation on Math Reasoning with MCTS-Q. We show the generations between CoLAA using MCTS-Q baseline
39

---

Controlling Large Language Model with Latent Action
using MCTS-Q and CoLA (using FTA-P).
Problem 1: In how many ways can 8 people sit around round table if 3 of the people -- Pierre, Rosa, and Thomas -- all want to sit together? (Iwo seatings are considered the same if one is a rotation of the other.) Answer: 720
CoLA with MCTS-Q: 1. Group Pierre, Rosa, and Thomas: Treat Pierre, Rosa, and Thomas as single entity. This reduces the problem to arranging entities (the group and the other 5 individuals) around the table.
2. Arrange the entities: There are 6! ways to arrange these entities around the table.
3. Arrange the individuals within the group: There 31 ways to arrange Pierre, Rosa, and Thomas within their group.
4. Account for rotations: Since the table is round, each arrangement can be rotated in different ways (one for each person as the starting point). Therefore, we need to divide by 6 to correct for these rotations.
5. Calculate the total number of arrangements: 720x6 720
6xx3
The final answer is 720
CoLA (SFT): 1. Calculate the total number of ways to arrange 8 people around a round table:
Since the table is round, we can fix one person's position to avoid counting rotations as different arrangements. This gives us 7! ways to arrange the remaining people.
2. Calculate the number of ways to arrange Pierre, Rosa, and Thomas
together:
Treat Pierre, Rosa, and Thomas as single unit. This reduces the problem to arranging 6 units (the unit of Pierre, Rosa, and Thomas, and the other 5 individuals) around the table. Fixing one of these units, we have 5! ways to arrange the remaining units.
3. Calculate the number ways to arrange Pierre, Rosa, and Thomas within their unit:
Pierre, Rosa, and Thomas can be arranged among themselves in 3! ways.
4. Combine the arrangements:
Multiply the number of ways to arrange the units by the number of ways to arrange Pierre, Rosa, and Thomas within their unit: x 120 6 720
5. Adjust for the round table:
Since the table is round, we need to divide by the number of people (8) to avoid overcounting rotations: 720 90
The final answer is
40

---

Controlling Large Language Model with Latent Action
Baseline with MCTS-Q We' 11 consider the three people who want to sit together single unit. That leaves us with 5 units to arrange around the table (the 3-person unit and the other 5 individuals).
1. Calculate the number of ways to arrange these 5 units around the table. Since they are indistinguishatt shable in terms of rotation, we use the formula n items, which is ()! Here, n n=5:
for circular permutations Ways to arrange 5 units (5 1)! 4! x 24
2. Now, within the 3 3-person unit, Pierre, Rosa, and Thomas can be arranged among themselves in 3! ways: =1
3. Multiply the number of ways to arrange the units by the number of ways to arrange the individuals within the unit to get the total number of arrangements:
Total arrangements = 24 x 6 = 144
144

---

Controlling Large Language Model with Latent Action
Problem 2: If 0.1331 is written fraction
fracab with and b being positive integers with greatest common divisor of 1, what ig 6+
Answer: 1030
CoLA with MCTS-Q: 1. Let 0.133I. Then, multiplying both sides by 10000
(since the repeating part has four digits), get:
10000z = 1331.1331
2. Subtract the original & from this new equation:
10000x - x = 1331.1331 0.1331
9999r = 1331
Solving for 2, we find: 1331
9000
4. To simplify 1331 find the greatest common divisor of 1331 and 9999.
9999
Since 9999 is 3 x
11 x 101 and 1331 is 113, the ged is 11. Thus, we can simplify the fraction:
1331 1331#11
121
9999 9099: = 909
5. The sum of the numerator and the denominator is:
121 + 909 1030
COLA (SFT): 1. Let x 0.133. To eliminate the repeating decimal, multiply
x by 10000 (since the repeating part has four digits):
10000r 1331.1331
2. Subtract the original from this equation:
10000z 133I - 0.133I
9999r 1331
3. Solve for 2:
1331
I 9999
4. Simplify the fraction 1331 by finding the greatest common divisor
9999 (GCD) of 1331 and 9999. Since 1331 is perfect square 113 and 9999 is not
multiple 11, the GCD is 1. Thus, the fraction is already in its simplest form.
5. Calculate + b where a= 1331 and b = 9999:
a a = 1331 + 9999 = 11330
Baseline with MCTS-Q Let 0.1331. Then 1000r 1331.1331, SO 1000z x 1331.1331 - 0.133I, which simplifies to 1000x - = 1331 Solving for H, we get 1331
999x = 1331, so x 999
To simplify 999 we find the greatest common divisor the numerator and the denominator. Since 999 3 x 37 and 1331 113, the greatest common divisor of 999 and 1331 is 1. Thus, the fraction is already in its simplest form.
Therefore, a+b b = 1331 + 999 2330
42


Reviews:
[{"review_id": "r2pe8zQ3sK", "reviewer": "Reviewer_iaGP", "summary": "This paper proposes to learn a more compact latent \"action\" space for pretrained LLMs to improve controllability. This latent action space is learned by an autoencoder where the encoder is an inverse dynamics model p(a_t|x_0:t+1), the decoder is a transition model p(x_t+1|x_0:t, a_t), the conditional prior is a policy p(a_t|x_0:t). The training objective for this latent action space is similar to VQ-VAE, but parametrize p(a_t|x_0:t+1) with Gumbel Softmax and the KL term in the ELBO is parametrized by the entropy of p(a_t|x_0:t+1). The authors further provide two SFT approaches based on finetuning either the inverse dynamics or the policy, as well as a RL objective on the latent policy. Experiments show that the proposed method can induce a latent space with only a slight degradation of the base pretrained model. The authors further show that RLHF and SFT over latent actions appear to be more effective than on the base model with tokens as actions.\n\n### After rebuttal ###\nThanks the authors for addressing my concerns. I will keep my initial accepting rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for helpful feedback and positive assessment of the novelty in our latent action control approach. We address the reviewer's concern below, with key claims highlighted in bold for clarity.\n\n**Q1:** In the ablation study, why did increasing the model size and the data size lead to worse performance? This doesn't seem to obey the scaling laws of LLMs. I hope the authors would like to provide some explanation.\n\n**A1:** Thank you for your question. We’d like to clarify that our findings **do not conflict with the scaling laws of LLMs**. The slight performance degradation observed in Figure 6 stems from the \"forgetting issue\" during continued pre-training of Llama-3.1-8B. This occurs because our experimental data differs from Meta’s original pre-training data, leading to partial knowledge forgetting. Similar effects have been documented in prior work [1]. To our knowledge, scaling laws [2] are primarily applied to large-scale pre-training, not continued training with limited data.  \n\nWe clarify that we show Figure 6 to indicate that the original Llama-3.1-8b is a proper baseline. We will revise the descriptions to modify the explanation of this experiment in Section 5.4 to explain it more clearly. We understand that the reviewer may expect improved performance with more data. While this could be possible by enhancing LAMP with additional insights and techniques from continued training, such an investigation is beyond the scope of this paper. We appreciate your feedback and consider this an important direction for future work.\n\n[1] Shao et al. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300 (2024).\n\n[2] Kaplan et al. Scaling Laws for Neural Language Models. arXiv:2001.08361 (2020).", "claims_and_evidence": "The thesis claim is that the latent action, due to its lower dimensionality, can be a more controllable alternatives to the token space. The authors provide evidence that SFT and RLHF on the latent space turn out to perform better.", "methods_and_evaluation_criteria": "The proposed method follows the standard framework of VQ_VAE. The adopted evaluation criteria is commonly used.", "theoretical_claims": "There is no theoretical claim.", "experimental_designs_or_analyses": "The authors first validate the success of latent space induction by showing that there is only a slight degrade of the performance after learning the latent space modules. \nThey then test how well the latent actions control the tokens generation. \nThe main results are the performance gain from SFT and RLHF on the latent action space. \nThe authors also provided ablation study for the increased token and parameter numbers. However, this appears to be skeptical to me. The reported ablation show degradation of the base model in terms of increased model size and data size, which is very counter-intuitive.", "relation_to_broader_scientific_literature": "The idea of using a more compact space for SFT and RLHF can improve the efficiency of the SOTA pipeline.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The idea of learning latent action space to improve controllability of LLM appears to be novel to me.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "In the ablation study, why did increasing the model size and the data size lead to worse performance? This doesn't seem to obey the scaling laws of LLMs. I hope the authors would like to provide some explanation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RUO8fdxRLt", "reviewer": "Reviewer_X62N", "summary": "The paper presents Latent Action governed world Model from Pre-trained LLM (LAMP), a novel reinforcement learning (RL) approach designed to control Large Language Models (LLMs) by learning a compact, latent action space rather than relying on conventional token-level actions. LAMP leverages an inverse dynamics model to create discrete latent actions and a policy model to guide these actions for downstream tasks. Experimental results on the Llama-3.1-8B model show moderate improvements in performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer's thoughtful feedback and positive recognition of the novelty of the perspective of RL from observation. We are grateful for the opportunity to address the reviewer's concerns and further clarify our methodology. \n\n**Q1:** While the claims regarding improved controllability is generally supported, the main concern is that the performance gains demonstrated by LAMP are relatively marginal compared to baseline methods. Specifically, the reported improvements such as Table 1 and Table 2, appear limited in practical significance given the complexity of the proposed method.\n\n**A1:** We apologize for the unclear expression. We clarify that **Tables 1 and 2 show only \"me too\" but not the main purpose of LAMP**, which is to enhance the downstream tasks using reinforcement learning.\n\n> Explanation: Table 1 shows the performance of LAMP model with behavior cloning policy, which should be called as LAMP-BC. We train LAMP-BC just to show that the LAMP structure can perform as an ordinary LLM, as presented results in Table 1. Further, Table 2 shows supervised fine-tuned performance of both the base LLM and LAMP-BC, which shows LAMP-BC can have a slightly better performance. We will revise the corresponding descriptions in Sections 5.1 and 5.2 to strengthen the presentation for the propose of these two experiments to explain this more clearly. \n\nTo provide an overview of the power of our approach, we summarize and explain our RL results below: \n\n- For preference alignment tasks, our method achieves superior win rates and robustness via RLHF across different KL settings. In **Figure 2**, our LAMP model achieves an average of 64% win rate at KL=0.01 and over 90% win rate at KL=0.00. In **Figure 3**, our LAMP model achieves a slight performance improvement under the KL=0.00 setting compared to the standard KL=0.01 configuration, indicating a more robust alignment of our LAMP model.\n- For mathematical reasoning, after fine-tuning and RL optimization, our method achieves 42.6 on MATH500, comparing with 38.2 of baseline. And achieves 68.2 on a more flexible MCTS method, comparing with 63.0 of baseline.\n\nThese results show that **our method achieves significant improvements by introducing latent actions for RL training**.\n\n**Q2:** Additionally, the claim about the semantic interpretability of latent actions is not fully convincing. It remains unclear how these high-level semantic categories meaningfully translate into effective control or actionable improvements.\n\n**A2:** We would like to clarify that **we did not claim the semantic interpretability of the extracted latent actions**. In fact, Figure 12 (Appendix C.2, Page 18) aims to show that latent actions have higher semantics than tokens. We will revise the claim in Section 4 to state that the connection of the latent actions to semantic is weakly correlated, avoiding cause misunderstandings.\n\nTo address your concern about the effectiveness of latent actions, we clarify that the higher-level latent actions not only enhance the reinforcement learning performance due to reduced latent action space, but also mitigate the reward hacking issue (Figure 3), which can lead to more stable reinforcement learning.", "claims_and_evidence": "While the claims regarding improved controllability is generally supported, the main concern is that the performance gains demonstrated by LAMP are relatively marginal compared to baseline methods. Specifically, the reported improvements such as Table 1 and Table 2, appear limited in practical significance given the complexity of the proposed method.\n\nAdditionally, the claim about the semantic interpretability of latent actions is not fully convincing. It remains unclear how these high-level semantic categories meaningfully translate into effective control or actionable improvements.", "methods_and_evaluation_criteria": "- The proposed methods make sense and the motivation is clear. \n- The connection between LLM training and observation-only RL is novel.\n- The evaluation criteria are reasonable.", "theoretical_claims": "The paper does not present theoretical claims.", "experimental_designs_or_analyses": "The experimental design and analyses appear generally sound.", "relation_to_broader_scientific_literature": "The paper situates itself effectively within the broader scientific literature, clearly outlining the limitations of token-level control and the motivation drawn from \"RL from observations\" literature. The contributions are well-contextualized relative to recent developments such as RLHF and latent-space modeling.", "essential_references_not_discussed": "This paper appears comprehensive in its coverage of related works and does not omit essential references critical to understanding the contributions.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wdfL00Ukba", "reviewer": "Reviewer_mZHY", "summary": "Existing LLMs often rely on token-level actions that may be overly large and inefficient. This paper proposes learning a compact and latent action space to improve controllability and exploration in RL. Specifically, the authors augment a latent action space with a pre-trained LLM to form a latent action language world model. This latent action model is extensively trained from the pre-training dataset and tuned in the post-training stage. Experiments with Llama-3.1-8B as the base model demonstrate that using this latent action model for RL training achieves better controllability on multiple tasks, including multiple preference alignment tasks, math reasoning task, as well as the improved and flexible searching over the token-level action framework.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer's valuable feedback and recognition of our latent action approach.\n\n**Q1:** Claims are not supported by clear evidence: Therefore, LAMP aims to decouple these semantics by assigning them to different and limited latent actions, enabling more efficient fine-tuning, alignment, and search by latent action control. ==> I did not see such evidence.\n\n**A1:** For claim 1, we realize that this claim may potentially lead to misunderstandings and would like to clarify. Figure 12 and Figure 10 (Appendix C.2) aim to show that latent actions have higher semantics than tokens. We will revise the statement in Section 4 to clarify that the connection between latent actions and semantics is weakly correlated, which should avoid misunderstandings.\n\n**Q2:** Claims are not supported by clear evidence: However, in our experiments, we observed that this suffers from codebook collapse, where only a limited number of actions are activated during training.\n\n**A2:** For claim 2, we compare VQ-VAE with our proposed method in Figure 8 (Appendix C.1) by counting the number of actions with non-zero usage. Results show that VQ-VAE suffers from severe codebook collapse, while our approach effectively maintains the utilization of codes in the codebook. We will add a reference to this claim to make it more supported.\n\n**Q3:** However, the authors should show that their methods are appliable to other LLMs besides Llama-3.1-8B.\n\n**A3:**  We have implemented the LAMP architecture with Qwen-2.5-math-1.5B (Appendix C.1, Figure 9). We observed similar behavioral patterns to those in Llama-3.1-8B, including stable loss reduction and maintained codebook diversity without collapse. This provides preliminary evidence for LAMP's generalizability. However, due to the limitation of computation resources, we did not proceed to full-scale training and subsequent downstream task optimization.\n\n**Q4:** In the generation process, we compute the next token from the world model in a greedy manner. We only randomly or greedily sample actions from the policy model. ==> what's your meaning?\n\n**A4:** This sentence describes LAMP's language generation process. During the RL training phase, LAMP's policy model explores by randomly sampling latent actions, while during the evaluation phase, it deterministically (greedily) samples latent actions to ensure consistent results by selecting the maximum probability action. We will further revise here to make a clearer description.\n\n**Q5:** Some writing issues, including typos and redundant descriptions.\n\n**A5:** We will improve the writing in the revision.\n\n**Q6:** line-199: why the current time embedding e^i_t has a superscript i?\n\n**A6:** Here we aim to differentiate the embedding from other model input/output embeddings, and $i$ represents the inverse dynamics model.\n\n**Q7:** I did not see the necessary of MCTS-Q? It makes the methods complex?\n\n**A7:** MCTS-Q demonstrates how our latent action formulation can effectively integrate techniques that were previously difficult to apply in token space. For instance, Q-learning-based uncertainty estimation is a key technique for pruning the search space in MCTS, but scaling it in token space is challenging. Thanks to our latent action space design, however, this approach now shows promising results. Importantly, MCTS-Q does not introduce significant complexity—instead, it makes MCTS more effective.\n\n**Q8:** The results did not show that the proposed LAMP are better than baselines with a clear margin.\n\n**A8:** We guess the reviewer draws this conclusion from the results in Table 1 and Table 2. We want to clarify the performance gains of our approach. Our method primarily demonstrates superior downstream performance in multiple tasks. **For preference alignment**, we achieve an average win rate of **64%** across four tasks and exceed **90%** under the KL=0.00 setting. **For math reasoning**, we also shows significant improvement (**42.6** vs 38.2). **For search**, combined MCTS with Q, our method can achieve better performance (**68.2** vs 63.0). **We also explore an agentic task called Alfworld**[1], which is an embodied multi-turn interactive task. With the same fine-tuning and RL scheme, results shown in following table demonstrate superior performance of LAMP on it:\n\n| Method | Seen | Unseen |\n|-|-|-|\n|Llama-3.1-8B-SFT|68.6|67.9|\n|Llama-3.1-8B-RL|68.6|71.6|\n|LAMP-FTA|75.7|70.9|\n|LAMP-RL|**77.9**|**74.6**|\nWe also need to clarify that Tables 1 and 2 show only \"me too\", which only shows the comparable initialized performance of LAMP. We will revise the descriptions in Section 5.1 and Section 5.2 to strengthen the presentation of these aspects. Please let us know if you have additional concerns about the performance gains of our approach.\n[1] Shridhar et al. ALFWorld: Aligning Text and Embodied Environments for Interactive Learning. arXiv:2010.03768 (2020).", "claims_and_evidence": "Most claims are supported by clear evidence. \n\nSome are not:  \n1) Therefore, LAMP aims to decouple these semantics by assigning them to different and limited latent actions, enabling more efffcient ffne-tuning, alignment, and search by latent action control.   ==> I did not see such evidence.     \n\n2) However, in our experiments, we observed that this suffers from codebook collapse, where only a limited number of actions are activated during training.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "Yes. The four questions investigated are important and the experimental designs are reasonable.  \n\nHowever, the authors should show that their methods are appliable to other LLMs besides Llama-3.1-8B.", "relation_to_broader_scientific_literature": "The effectiveness of latent action will be related to most of the LLM researchers. Especially those have RL background.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:  \nThe latent action is a very interesting topic, and I think this is important for the future of LLMs  \nThe methods are reasonable, and effective.    \nThe paper is generally well-written, easy-understanding.\n\nWeaknesses:  \nsmall writing issues, see [Other Comments Or Suggestions]    \nI did not see the necessary of MCTS-Q? It makes the methods complex?   \nThe results did not show that the proposed LAMP are better than baselines with a clear margin.", "other_comments_or_suggestions": "‘’RL from observations only‘’ ==> ``RL from observations only”   \nIn this section, we verify our framework across multiple training stages. First, we introduce the basic setting of our model design. Then we validate LAMP in multiple training stages.  ==> is redundant  \nIn the generation process, we compute the next token from  the world model in a greedy manner. We only randomly or greedily sample actions from the policy model.  ==> what's your meaning?", "questions_for_authors": "line-199: why the current time embedding e^i_t has a superscript i?    \nI did not see the necessary of MCTS-Q? It makes the methods complex?    \nThe results did not show that the proposed LAMP are better than baselines with a clear margin.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel approach - connecting LLMs to learning from observation - and uses latent actions to steer a language model. This is a novel idea, and may be of interest to the community, especially given the number of very similar ideas that are now being pursued in the era of LLMs. That being said, experiments remain narrow with only Llama 3.1 8B, and it remains unclear how applicable the approach is to more open ended problems.
Decision: Accept (poster)