Language Model Adaption for Reinforcement Learning with Natural Language Action Space (2024)

Abstract: Reinforcement learning with natural language action space often suffers from the curse of dimensionality due to the combinatorial nature of the natural language. Previous research leverages pretrained language models to capture action semantics and reduce the size of the action space. However, since pretrained models are typically trained on general corpora, there can be an unpredictable mismatch between the priors encoded in pretrained models and the characteristics of the specific RL environment. To address this issue, we propose Mutual-Information Regularized Policy Optimization, MIPO. MIPO enables implicit and dynamic reduction of the action space. Starting from the prior provided by the pretrained language model, our method dynamically adjusts the prior during the learning process based on the guidance of mutual information regularization. Theoretically, we demonstrate that this policy optimization process leads to the monotonic improvement on the mutual-information regularized RL objective. Empirically, we conduct experiments in various environments and demonstrate the effectiveness of MIPO.

heading: 1 Introduction
text: Deep reinforcement learning (RL) has gained great success from mastering the game of Go (Silver et al., 2017) to training state-of-the-art large language models (Ouyang et al., 2022). Despite its success in various domains, the low sample efficiency of RL has consistently hindered its wide adoption in real-world applications where the cost of interaction is not negligible (Li, 2017). Furthermore, as suggested in Azar et al. (2017), the sample complexity of RL algorithms will be inevitably influenced by the size of the action space. This issue becomes even more pronounced in scenarios with enormous action spaces, such as multi-agent rein-
†Corresponding author.
forcement learning (Zhang et al., 2021) and openworld reinforcement learning (Yuan et al., 2023a), exacerbating the problem of sample efficiency.
One natural solution to tackle the challenges of large action space is to eliminate irrelevant actions accordingly to avoid the exploration of unnecessary state-action pairs. While a general principle of action space reduction is hard to find, for reinforcement learning problems with natural language action space (He et al., 2016), leveraging action semantics provides a means to infer the potential usefulness of actions. Previous work has explored the use of pretrained language models for action space reduction (Martin et al., 2022). Despite being a viable solution, pretrained models encode prior knowledge derived from general corpora, which may differ from the specifics of the actual environment. To alleviate such a mismatch, some studies (Yao et al., 2020; Xu et al., 2022) utilize the domain-specific dataset to finetune the pretrained model for better adaptation. Others (Zahavy et al., 2018; Ammanabrolu and Hausknecht, 2019; Jain et al., 2020) assume the feedback of action admissibility from the environment, which assesses the relevance of each action to the current task, and use it to prune the action space. However, neither the domain-specific dataset nor the admissibility signal is always available. This leads to the question of whether it is possible to adjust the prior knowledge without incurring additional costs or making assumptions.
To address this challenge, we propose MutualInformation Regularized Policy Optimization, namely MIPO. Unlike previous work where the hard action mask is used to prune the action space, we use the pretrained language model to generate a prior policy over the action space of the agent. The agent is subsequently trained to optimize the environmental reward as well as being close to this prior. Then, based on the marginalization of the agent’s policy, we dynamically adjust the general
1620
prior provided by the pretrained language model to adapt to the specific environment. Such a learning process can be unified as a mutual-information (MI) regularized reinforcement learning problem, where the standard RL objective is augmented with the mutual information between the agent’s policy and non-text state information. In theory, we prove that MIPO achieves the monotonic improvement on the MI-regularized objective. To validate our method in practice, we perform a series of experiments on text-based games. The experimental results show that MIPO outperforms baselines in terms of training speed and final performance.
Our contribution can be summarized as follows. First, we propose MIPO, an MI-regularized policy optimization method that achieves implicit dynamic action space reduction without requiring additional datasets or assumptions. Then, we prove that our method achieves monotonic improvement in the MI-regularized objective. Finally, we empirically validate our method and demonstrate its effectiveness.

heading: 2 Related Work

heading: 2.1 RL in Text-Based Games
text: In text-based games (Côté et al., 2019), the environment takes an action described in natural language and returns the natural language description of the current state. Many studies have focused on solving text-based games using only text input. While early work among them mostly focuses on the network structure (Narasimhan et al., 2015; He et al., 2016), recent studies incorporate techniques such as hierarchical structure (Adolphs and Hofmann, 2020; Zhu et al., 2023), multi-passage reading comprehension (Guo et al., 2020), and well-designed exploitation and exploration mechanism (Tuyls et al., 2022) to achieve better performance.
Other than using pure text input, another line of research combines the knowledge graph into the learning algorithm to achieve better state representation. For example, Ammanabrolu and Riedl (2019); Ammanabrolu and Hausknecht (2019); Adhikari et al. (2020) consider different GNN structures for the representation of knowledge graph, Xu et al. (2020, 2021) introduce a hierarchical structure for the state representation and the decisionmaking process, and Atzeni et al. (2021) combine case-based reasoning with knowledge graph for the better exploration strategy.
Technically, our algorithm can be combined with
any of the above methods. However, we choose to combine the knowledge graph into our algorithm, as it is a more general setting where the state does not contain only text information.

heading: 2.2 Action Space Reduction
text: The sample complexity of the RL algorithm generally grows with the size of the action space. This becomes a problem for text-based games due to its combinatorial nature on the action space. However, as natural language actions can be understood by humans and language models, various methods have been proposed to reduce the size of the action space using its semantic nature. In Ammanabrolu and Riedl (2019), a rule-based method is used to extract information from the knowledge graph and perform subsequent action selection. While this method requires domain knowledge about the environment and the corresponding knowledge graph, some other methods (Zahavy et al., 2018; Ammanabrolu and Hausknecht, 2019; Jain et al., 2020) assume the feedback of action admissibility from the environment and use it to prune the action space. To remove the assumption about the environment or the requirement of domain knowledge, Martin et al. (2022) utilize the power of pretrained language models to perform the action space reduction. However, as prior knowledge of the pretrained language model is generally derived from general corpora, it may exhibit a certain level of mismatch of the actual environment. To alleviate such a mismatch, domain-specific dataset (Yao et al., 2020; Xu et al., 2022) is collected to further tune the pretrained language models for better adaption.
Unlike these works, our method makes no assumption on the domain knowledge, the admissibility feedback, or the domain-specific dataset. The adaption of the pretrained language model is guided by the policy of the agent, which is derived from the mutual-information regularized objective. To the best of our knowledge, only Shi et al. (2023) uses the same setting of our paper, where the pretrained language model is tuned via self-supervised learning. We take it as a baseline in our experiments.

heading: 2.3 Regularization in RL
text: The goal of RL algorithms is to learn a policy that optimizes a single objective, the cumulative discounted reward. However, in practice, we may want to also incorporate other features into the learned policy. This is usually achieved via in-
corporating a regularization term into the policy optimization process. For example, SQL (Haarnoja et al., 2017) and SAC (Haarnoja et al., 2018a) use the entropy of policy as a regularization term to encourage exploration. The entropy term can be considered as the KL divergence between the policy and a uniform prior policy, therefore we can change the uniform prior to other fixed prior to achieve different purposes. For example, in offline RL (Levine et al., 2020), BRAC (Wu et al., 2019) and CQL (Kumar et al., 2020) use the KL to control the distance between the learned policy and the behavior policy to mitigate the instability arising from previously unseen actions. Such a regularization is also used in PPO (Schulman et al., 2017) and is further used for the optimization of large language models (Ouyang et al., 2022; Rafailov et al., 2023). While previous methods consider a fixed prior, Grau-Moya et al. (2018) and Leibfried and Grau-Moya (2020) consider a dynamic prior learned along the policy optimization process as an improvement of SAC (Haarnoja et al., 2018a). Unlike these methods, we consider the case where the state information can be decomposed into textual and non-textual components, we then use a pretrained language model as the prior to handle the textual components and adapt this language model dynamically during the policy optimization process.

heading: 2.4 Foundation Models for RL
text: Foundation models (Yang et al., 2023) pretrained on diverse and large-scale datasets contain valuable prior knowledge, providing useful insight on the downstream tasks. Given the success of foundation models, using foundation models to achieve a better decision-making process has become a popular research topic. Except from the action space reduction, foundation models have also been used in many different perspectives to enhance the performance of RL algorithms. For example, it can serve as a reward model (Kwon et al., 2022), dynamic model (Zhong et al., 2022), or affordance function (Ahn et al., 2022). It can also be used to generate intrinsic reward (Du et al., 2023) to encourage exploration. For long-horizon tasks, foundation models have been used for different hierarchical methods (Nottingham et al., 2023; Wang et al., 2023; Yuan et al., 2023b) to generate sub-goals. It can also be used to parameterize the agent’s policy directly (Carta et al., 2023). Unlike these methods, we focus on using pretrained language model to
reduce the action space.

heading: 3 Background
text: Reinforcement Learning Problem. We often formulate the RL problem as a Markov Decision Process (MDP) (Bellman, 1957). An MDP can be defined by a tuple ⟨S,A, P,R, ρ0, γ⟩. At each state s ∈ S, the agent can choose one of the actions a ∈ A to execute and receive a reward signal r(s, a) from reward function R(s, a) : S ×A → R. The state will then transition from the current state s to the next state s′, governed by transition function P (s′|s, a) : S ×A×S → [0, 1]. In this problem, the goal of the agent is to properly select action following policy π(a|s) at each state to maximize the expected cumulative discounted reward Eρ0,π,P [ ∑∞ t=0 γ
tr(st, at)], where the expectation is over initial state distribution ρ0, agent policy π and transition function P , γ ∈ (0, 1) is the discount factor, and t denotes timestep.
The RL objective can also be represented by the reward function and the state marginal distribution as follows: (1− γ)Eρ0,π,P [ ∞∑
t=0
γtr(st, at) ] = Eρπ(s) Eπ(a|s) [ r(s, a) ] ,
where ρπ(s) = (1− γ) ∑∞ t=0 γ tP πt (s), and P π t (s) denotes the probability that policy π visits state s at timestep t. RL with Natural Language Action Space. For RL with natural language action space, each action a is paired with the corresponding semantic and therefore can be encoded into the action representation ha using pretrained language models. Without loss of generality, we further assume each state can be decomposed into two components, s = {stext, snon-text}, where stext represents the state information that can be aligned with action semantics and snon-text represents the state information that cannot be aligned with action semantics. We choose pretrained language models (Sun et al., 2022) for the state-action semantic alignment, as in our experiment scenario, the state information can be decomposed into knowledge graph information and text information. GATA. Our implementation is based on GATA (Adhikari et al., 2020), which also serves as an important baseline in our experiments. In GATA, the state information is assumed to be decomposed into knowledge graph information sgraph and text information stext. Each component of the state information is encoded separately and subsequently
concatenated to form the joint representation of the state hs = [hsgraph , hstext ]. The state representation hs and the action representation ha will be then utilized as inputs of state-action value function Q(s, a), and this function will be updated via the following DQN (Silver et al., 2017) update rule to get the optimal state-action value function Q∗(s, a):
Q(s, a) = r(s, a) + γ EP (·|s,a) [ argmaxa′ Q(s ′, a′) ] .

heading: 4 Method
text: In this section, we present Mutual-Information Regularized Policy Optimization (MIPO). MIPO leverages a pretrained language model conditioned on the state variable stext to generate a prior policy πprior(a| stext). Then in Section 4.1, we incorporate the KL divergence between the current policy π(a|s) and the prior policy into the policy iteration process. In this way, we encourage the policy to optimize the environmental reward, as well as being close to the prior policy. Inspired by the mutual-information (MI) regularization, we exploit the marginalization of the agent’s policy to adapt the prior policy to the environment. As shown in Section 4.2, this alternative optimization process yields a monotonic improvement on the MI-regularized RL objective. Finally, in Section 4.3, we show how to implement this optimization process in practice.

heading: 4.1 KL-Regularized RL for Fixed Prior
text: Taking stext and the semantics of each action as inputs, a pretrained language model is able to produce a prior policy on the action space πprior(a| stext). Instead of directly using it as a hard mask for explicit action space reduction, we aim to learn a policy that can maximize the environmental reward as well as minimize the distance between the agent’s policy and the prior policy, which serves as implicit action space reduction.
We choose the KL divergence to measure the distance such that the policy evaluation operator can be defined as follows:
Q(s, a) = r(s, a)
+ γ EP,π [ Q(s′, a′)− α log π(a ′|s′) πprior(a′| stext′) ] ,
(1)
where the expectation is over P (·|s, a) and π(·|s′), and α is the coefficient for the trade-off between the actual return and the KL term. Note that if we take
π(·|s′) inside the expectation, the log term becomes DKL(π∥πprior). Then, the following optimization problem is used to achieve the policy improvement: π(·|s) = argmaxπ′ Eπ′ [ Q(s, ·)− α log π
′(·|s) πprior(·| stext) ] .
(2)
As shown in Su and Lu (2022), by repeatedly applying Equation (1) and (2), the policy converges to the optimal policy in terms of the following the KL-regularized RL objective:
Eρ0,π,P [ ∞∑
t=0
γt ( r(st, at)− α log π(at|st) πprior(at| stext,t) )] ,
which can alternatively be represented as:
Eρπ(s) Eπ [ r(s, a)− α log π(a|s) πprior(a| stext) ] .
The first term is for the reward maximization and the second term is for the KL distance minimization. It is worth noting that the regularization terms in Equation (1) and (2) are not duplicated, as we need to consider the effect of regularization in both the policy evaluation process and the policy improvement process to ensure the improvement over the regularized RL objective.

heading: 4.2 MI-Regularized RL for Dynamic Prior
text: In the previous section, we discuss incorporating the KL divergence between the policy and the prior into the RL objective. However, as the encoded knowledge in pretrained language models often stems from general corpora, there is a possibility that the prior provided by the pretrained language model does not align well with specific environments. Consequently, rigidly constraining the policy to be close to the prior may be suboptimal. To solve this problem, a mechanism is needed to dynamically adjust the prior to comprehend the environment.
As we make no assumption on the environmental feedback, the agent’s policy emerges as the sole teacher within this environment. Fortunately, it serves as a good one as it is designed to optimize the environmental reward, which contains enough information to differentiate between good and bad actions. Considering that the agent’s policy incorporates more inputs compared to the prior policy, we use the conditional marginalization of the agent’s policy to adapt the prior policy as follows:
πprior(a| stext) =Eρπ(snon-text | stext) [ π(a| stext, snon-text) ] , (3)
where ρπ(s) = ρπ(stext)ρπ(snon-text | stext) is the conditional factorization of the state marginal distribution of policy π. Such an association with the agent’s policy indicates that the adaptation of the prior policy should be performed on the states that are currently visited by the agent’s policy, which is crucial from a theoretical perspective as shown later.
Incorporating this prior adaptation process into the policy iteration framework leads to an alternating optimization process of (π, πprior). In the following theorem, we prove that this alternating optimization can be unified as the optimization of MI-regularized RL objective and ensures the monotonic improvement on this objective.
Theorem 1. If a sequence (πk, πpriork ) ∞ k=0 is obtained by iteratively applying Equation (1),(2) and (3), then it exhibits the monotonic improvement property on the MI-regularized RL objective, J(πk+1, π prior k+1 ) ≥ J(πk, π prior k ). Here, the MIregularized RL objective J(π, πprior) is defined as:
J(π, πprior) = Eρπ(s) Eπ [ r(s, a)− α log π(a|s) πprior(a| stext) ] .
Proof. See Appendix A.1.
In this theorem, we first treat the prior policy as a fixed one and prove that, by applying Equations (1) and (2), the policy iteration with a fixed prior leads to the improvement on the MI-regularized RL objective, J(πk+1, π prior k ) ≥ J(πk, π prior k ). Then, we fix the agent’s policy and perform the prior policy adaption by applying Equation (3), which also leads to the improvement on the objective, J(πk+1, π prior k+1 ) ≥ J(πk+1, π prior k ). Put them together, we can have the monotonic improvement J(πk+1, π prior k+1 ) ≥ J(πk, π prior k ).
In contrast to the policy iteration that optimizes the KL-regularized RL objective, we now have two
variables π and πprior. Therefore, the augmented term does not simply stand for the KL divergence between π and an arbitrary πprior, but it now stands for the mutual information between a and snon-text, MI(a; snon-text | stext).
From an information theoretical perspective, this objective implies that among policies sharing the same RL objective, we favor those with lower dependency on snon-text when stext is given. Such policies are easier to learn by the prior, as the prior takes only stext as the input.

heading: 4.3 MIPO framework
text: In Section 4.1 and 4.2, we glance at our algorithm from a theoretical perspective. In this section, we discuss the practical implementation of our algorithm using neural networks.
In MIPO, the agent is equipped with three different network modules. One critic networkQ(s, a; θ) takes the knowledge graph, text information, and the semantics of actions as input and outputs the state-action value for each action. One policy network π(a|s;ψ) also takes the knowledge graph, text information, and the semantics of actions as input and outputs the probability for each action. One prior network πprior(a| stext;ϕ) takes only the text information and the semantics of actions and outputs the prior probability for each action. The framework of MIPO is illustrated in Figure 1.
For the update of Q(s, a; θ), we follow Equation (1), and update it by minimizing the following TD error:
L(θ) = ED [( Q(s, a; θ)− [ r(s, a) + γ ( Q̂(s′, a′)
− α log π(a ′|s′) π̂prior(a′| stext′) )])2] ,
(4)
whereD denotes the replay buffer, and Q̂ and π̂prior are respectively the target critic and the target prior
policy. These target networks are used to stabilize the training of Q(s, a; θ) and periodically updated via copying the parameters of Q(s, a; θ) and πprior(a| stext;ϕ).
To update π(a|s;ψ), we follow Equation (2), and update it via minimizing the following objective: L(ψ) = ED [ Eπ(a|s;ψ) [ α log
π(a|s;ψ) π̂prior(a| stext)
−Q(s, a) ]] ,
(5)
where, again, the target prior π̂prior is used to stabilize the training of π(a|s;ψ).
Following Equation (3), we update the prior policy πprior(a| stext;ϕ) via the the maximum likelihood estimation as follows:
L(ϕ) = ED̃ [ Eπ(a|s) [ log πprior(a| stext;ϕ) ]] , (6)
where D̃ is another replay buffer much smaller than D, such that it stores only the recent transitions and properly approximates the state marginal distribution ρπ(s). The overall learning algorithm of MIPO is summarized in Appendix A.2.

heading: 5 Experiments

heading: 5.1 Experiment Settings
text: We conduct two sets of experiments on cooking games provided in TextWorld (Côté et al., 2019) and house-holding tasks based on VirtualHome (Puig et al., 2018).
For the TextWorld environment, all tasks are generated using the code provided in GATA (Adhikari et al., 2020) and follow the same task setup as described in GATA. In this experiment, agents will receive recipes described in natural language including ingredients and directions. It will also receive a knowledge graph including triplets like [white onion, fridge, in] describing the current state of the environment. Based on this information, agents will take natural language actions like "take white onion from fridge" to finish the given recipes. There are four difficulty levels, each containing 100 different tasks for training and 20 different tasks for test. In each task, the goal of the agent is to collect different ingredients and cook them based on the recipe.
We make the following modifications to make tasks more challenging. First, we remove the recipe information from the knowledge graph. Therefore, the agent has to use both the information from the
knowledge graph and the information from texts to make proper decisions. Second, some actions are filtered in a hard-code manner in GATA, we recover these actions and result in a larger action space. Third, the text information contains only recipe information. With the above modifications, our setting is much harder than the original setting in GATA. More details about experiments and hyperparameters are included in Appendix A.3.
For the VirtualHome environment, we customize a set of house-holding tasks where agents are supposed to follow the given natural language instructions to tidy the room. For example, agents will receive natural language instructions like "put waterglass-283 into / on kitchencabinet-238", and a knowledge graph including triplets like [waterglass-283, kitchentable-232, on]. Based on these information, agents will take natural language actions like "walktowards waterglass-283" to finish the given instructions. In difficulty 1, agents will be given one instruction, and in difficulty 2, agents will be given two instructions. For the knowledge graph used in the VirtualHome environment, we use a rule-based filter to remove task-irrelevant objects in the scenarios (for example, floor, wall, ceiling) and only keep task-relevant objects. We then consider task-relevant objects in the scenarios and their relations as edges. The relations we considered in the knowledge graph of VirtualHome include "inside", "on", "facing", "close", "is", and "holds".

heading: 5.2 Baselines
text: We consider the following three baselines compared against our method in the experiments.
GATA (Adhikari et al., 2020): a knowledge-
graph based method, where the relational graph convolution network (Schlichtkrull et al., 2018) is used to process the information of knowledge graph and a transformer (Vaswani et al., 2017) encoder is used to process the text information of both the state and the action. Putting them together, we can have the representation of each state-action pair and then use it as the input of Q(s, a). The agent interacts with the environment using an ϵgreedy policy and is trained via DQN (Silver et al., 2017). More specifically, we use the GATA-GTF variant for our experiment, as the construction of the knowledge graph is beyond the scope of this paper.
GATA-TrufLL (Martin et al., 2022): a reinforcement learning method paired with action space reduction using pretrained language models. As TrufLL was originally proposed for the language generation tasks, we re-implement TrufLL based on the framework of GATA. For the action space reduction, we use the probability threshold method as the truncation function. The truncation function is shown as follows:
gpth(α)(a| stext) = 1fLM(a| stext)>λ,
which keeps actions with a probability fLM(a| stext) greater than λ. We use a fixed language model to perform action space reduction, and the agent is trained as GATA with the reduced action space.
GATA-CSM (Shi et al., 2023): a reinforcement learning method paired with action space reduction using pretrained language models. To adapt the pretrained model to the current task, CSM collects trajectories with good performance using a heapbased replay buffer and adapts the pretrained model based on this replay buffer. As CSM was originally proposed for the action generation setting, we reimplement CSM based on the framework of GATA. We use the same action space reduction method and training method as GATA-TrufLL for CSM, except the language model is not fixed, but dynamically adjusted based on the collected trajectories.
GATA-SayCan (Shi et al., 2023): a postprocessing method to utilize the knowledge in the pretrained language model. In SayCan, we assume a learned state-action value function. When executed in the environment, the action is chosen based on not just the state-action value function, but also taking the probability provided by the language model into consideration as follows:
Qcombined(s, a) = Q(s, a)× fLM(a| stext).
The state-action value function used for SayCan is trained via GATA.
MIPO: the proposed method in this paper. We implement our algorithm based on GATA with minimum changes to adapt to the policy iteration. For the prior policy πprior, we use the output of the language model as the logit, and get the prior policy as follows:
πprior(a| stext) = exp(fLM(a| stext)/β)∑ a∈A exp(fLM(a| stext)/β) .

heading: 5.3 Implementation Details
text: We use DistilBERT (Sanh et al., 2019) for methods paired with a language model. To compute fLM(a| stext), we first get the last hidden state from the language model for both the text state information and the action semantic as the representation of the state and action. Then these representations are used to compute the cosine similarity, which serves as the probability fLM(a| stext). For methods that adjust the language model, we fix all previous transformer blocks and only train the last one.

heading: 5.4 TextWorld
text: Main Result. We plot the learning curves of all methods in Figure 2 and present all final results in
Table 1. MIPO obtains the best performance in all four difficulty levels. It achieves the improvement of 7.3%, 4.9%, 29.1% and 13.9% respectively in four different levels compared to the best baseline.
Except for the fourth difficulty level, GATATrufLL achieves a worse result compared to GATA, which indicates that the prior knowledge encoded in the language model deviates from the fact in the environment at least in some cases. As the language model is fixed in GATA-TrufLL, it has no chance to correct the mistake. For GATA-CSM, as it dynamically adjusts the language model based on the interaction with the environment, it always outperforms or achieves the same level of performance compared to GATA-TrufLL. However, it cannot consistently achieve better performance than GATA. One potential reason is that GATA-CSM uses a heap-based replay buffer to collect trajectories for language model adaption. Though this mechanism certainly collects high-quality data, those trajectories may be largely off-policy. Therefore, the language model may have known the fact of some high-quality states, but the agent may be not able to reach those states, which reduces the effectiveness of the action space reduction. On the other hand, our method collects trajectories from ρπ(s), which leads to an on-policy language model adaption. As the adaption is performed on-policy, it alleviates the mismatch on those states that are currently visited by the agent, therefore achieving effective action space reduction and outperforming all other baselines in all four difficulty levels. As we can see from Table 1, GATA-SayCan fails hard in all difficulty levels and is outperformed by its base policy GATA, which indicates that a mismatched language model can make things even worse when not properly used and further demonstrates the importance of the language model adaption. Generalization. We present the generalization ability of each algorithm in Table 2 (the performance on the test set for all difficulty levels). As we can see, MIPO outperforms all baselines in the four difficulty levels, though we do not make any claim about MIPO having the advantage in terms of gen-
eralization ability. This improvement in the test set may be simply the result of MIPO’s better learning in the training set. Note that GATA-CSM and GATA-TrufLL still rely on the language model for execution, making it challenging to deploy them on devices with limited computational resources. In contrast, MIPO distills the knowledge from the language model into the policy, resulting in significantly reduced computational requirements compared to GATA-TrufLL and GATA-CSM. This advantage of our algorithm highlights its suitability for environments with constrained computation resources. Ablation Study. To further verify our method, we perform an ablation study in this section. We propose an ablation baseline, GATA-KL. In GATAKL, all other perspectives are the same as MIPO, except the language model is fixed and will not be trained. As we can see from Figure 3, MIPO outperforms GATA-KL in all four difficulty levels, indicating that the performance improvement is indeed achieved by the dynamic adaption of the language model. As we can see, MIPO achieves the largest gap against GATA-KL in Difficulty 3. This is because the prior policy resulting from the language model has a relatively low performance on Difficulty 3 (it achieves 0.47 in Difficulty 3, 1.02, 1.45, and 2.15 in Difficulty 1, 2, and 4 respectively). Therefore the necessity of language model adaptation is demonstrated best in this difficulty level.

heading: 5.5 VirtualHome
text: We further evaluate our method in the VirtualHome environment. All final results are presented in Table 3. As we can see, the comparison among GATA, GATA-TrufLL, and GATA-CSM is similar to the TextWorld experiment, where GATATrufLL achieves a worse result compared to GATA and GATA-CSM outperforms or achieves the same level of performance compared to GATA-TrufLL but can not consistently outperform GATA. Such a similar result strengthens the conclusion we made in the TextWorld experiment and again highlights not only the necessity of the language model adaption but also the necessity of adapting the language model in the right way.
Different from the TextWorld experiment, GATA-SayCan achieves a much better result here. This is because that the language instruction in this experiment is more direct and clear, so the performance of the language model becomes bet-
ter than the previous experiments and can therefore improve the performance of its base model (GATA). However, we notice that GATA-TrufLL and GATA-CSM do not enjoy the same benefits. This is because it is challenging to set a proper threshold for the truncation function. For example, when we take "put waterglass-283 into / on kitchencabinet-238" as the instruction, for actions "walktowards kitchencabinet-235", "walktowards kitchencabinet-236", "walktowards kitchencabinet237" and "walktowards kitchencabinet-238", they all have a high probability with the given instruction (0.90 for "walktowards kitchencabinet-238" and 0.89 for others). But for another necessary action "walktowards waterglass-283", it only has a probability of 0.80. Therefore, to make sure we retain all necessary actions, we have to set a relatively low threshold (0.75) in this experiment, such that the truncated action space is still very large. For GATA-TrufLL and GATA-CSM, as they need to learn from scratch, it is still very challenging for them to achieve good performance with such a large action space. For GATA-SayCan, it only needs to make sure the probability of "walktowards kitchencabinet-238" is greater than other "walktowards kitchencabinet" actions (0.90 > 0.89), therefore it can fully enjoy the benefit of improved performance of the language model.
As for our method, it achieves the best performance in both difficulties, which matches the result in TextWorld experiments, and again demonstrates the effectiveness of MIPO.

heading: 6 Conclusion
text: In this paper, we propose MIPO, which achieves dynamic language model adaption and implicit action space reduction for reinforcement learning with natural language action space. We show the problem can be formulated as the optimization of the MI-regularized RL objective, and we theoretically prove that MIPO achieves monotonic improvement on the MI-regularized RL objective. We perform experiments on text-based games, and our MIPO outperforms all baselines, which justifies its effec-
tiveness.
Limitations and Potential Risks
One limitation of this work is the convergence in terms of policy. As suggested by the theory, we can only guarantee the monotonic improvement on the objective, but the policy may oscillate between several equivalent points. Although we do not observe any significant impact in our experiment, it is still less satisfying and will be addressed in our future work. One another limitation of our work is the value of α, which is currently set heuristically. Although we find it easy to tune in our experiment, a fully automatic method, like the one used in SAC (Haarnoja et al., 2018b), to set α will be more welcomed and will be included in our future work. To the best of our knowledge, there is no obvious risk to our work.

heading: Acknowledgements
text: This work was supported in part by NSF China under grant 62250068.

heading: A Appendix
text: A.1 Proof of Theorem 1 Theorem 1. If a sequence (πk, πpriork ) ∞ k=0 is obtained by iteratively applying Equation (1),(2) and (3), then it exhibits the monotonic improvement property on the MI-regularized RL objective, J(πk+1, π prior k+1 ) ≥ J(πk, π prior k ). Here, the MI-regularized RL objective J(π, π prior) is defined as:
J(π, πprior) = Eρπ(s) Eπ [ r(s, a)− α log π(a|s) πprior(a| stext) ] .
Proof. We begin our proof by the following definition of the KL-augmented reward:
rππprior(s, a) = r(s, a)− αEs′∼P,a′∼π [ log π(a′|s′) πprior(a′|s′) ] .
To avoid the discussion on the corner case, we assume π and πprior take the form of ϵ-soft policy (Sutton and Barto, 2018).
Based on this definition of the KL-augmented reward, and the following way of the policy evaluation (7) and the policy improvement (8):
Q(s, a) = r(s, a) + γ EP,π [ Q(s′, a′)− α log π(a ′|s′) πprior(a′| stext′) ] , (7)
π(·|s) = argmaxπ′ Eπ′ [ Q(s, ·)− α log π ′(·|s) πprior(·| stext) ] , (8)
for a pair of (πk, π prior k ), we can follow Theorem 2 in Su and Lu (2022) to show the following conclusion:
J(πk+1, π prior k ) ≥ J(πk, π prior k ).
For the new policy πk+1 and the old prior policy π prior k , we have the following objective:
J(πk+1, π prior k ) = Eρπk+1 (s) Eπk+1(a|s)
[ r(s, a)− α log πk+1(a|s)
πpriork (a| stext)
] .
To update πpriork with a fixed πk+1, maximizing the above objective is equivalent to minimizing the following objective:
Ĵ(πk+1, π prior k ) = Eρπk+1 (s) Eπk+1(a|s)
[ log
πk+1(a|s) πpriork (a| stext)
] .
Therefore, for πpriork+1 that takes the following form of the prior adaptation:
πpriork+1 (a| stext) = Eρπk+1 (snon-text | stext) [ πk+1(a| stext, snon-text) ] ,
similar to Lemma 10.8.1 in Cover (1999), we can have the following proof:
Ĵ(πk+1, π prior k )− Ĵ(πk+1, π prior k+1 )
= Eρπk+1 ,πk+1 [ log πk+1(a|s) πpriork (a| stext) − log πk+1(a|s) πpriork+1 (a| stext) ]
= Eρπk+1 (stext) [ Eρπk+1 (snon-text | stext) Eπk+1 [ log
πpriork+1 (a| stext) πpriork (a| stext)
]]
= Eρπk+1 (stext) [ E πpriork+1 [ log πpriork+1 (a| stext) πpriork (a| stext) ]]
= Eρπk+1 (stext) [ KL(πpriork+1 ||π prior k ) ]
≥ 0
such that:
Ĵ(πk+1, π prior k+1 ) ≤ Ĵ(πk+1, π prior k ),
which indicates:
J(πk+1, π prior k+1 ) ≥ J(πk+1, π prior k ) ≥ J(πk, π prior k ),
and concludes our proof.
A.2 Pseudocode
Algorithm 1 MIPO 1: Initialize θ and ψ for critic network and policy network, load parameters ϕ for prior network; 2: Initialize target critic and target prior policy using θ and ϕ; 3: Initialize replay buffer D and | D̃ | (≪ |D |); 4: for i = 1 to I do 5: s← s0 6: for t = 0 to T do 7: Sample action at from π(at|st;ψ) 8: Execute action at and get rt, st+1 9: Store (st, at, rt, st+1) in D and D̃
10: Sample batch B from D 11: Update Q(s, a; θ) with B by (4) 12: Update π(a|s;ψ) with B by (5) 13: Sample batch B̃ from D̃ 14: Update πprior(a| stext;ϕ) with B̃ by (6) 15: end for 16: if i mod update_interval = 0 then 17: Update target critic and target prior policy with θ and ϕ 18: end if 19: end for
A.3 Experiment details For TextWorld experiments, our method and all baselines are implemented using PyTorch based on the code of GATA (Adhikari et al., 2020)1. We use the default hyperparameters of GATA-GTF for all methods, except the maximum number of episodes is set from 100k to 60k and we do not use any pretrained word embedding and set all word embeddings to be trainable. As MIPO uses an actor-critic framework, we add an actor head that contains 3 fully-connected layers with 64/64/1 units activated by ReLU except the last one. This actor head takes the state-action representation provided by the GATA network with gradient detached and produces the logit for each action. The learning rate for the policy network and prior network is 5 × 10−4. Except for the above-mentioned modifications, every detail is set to be the same as GATA. For VirtualHome2 experiments, the maximum number of episodes is further reduced to 20k to accelerate the training process. Some hyperparameters are also changed to adapt to this faster training process, for example, the batch size is set to 128 and the multi-step hyperparameter is set to 1.
For MIPO3 and GATA-KL, two hyperparameters α, β are added. β is fixed as 0.04 for all difficulty levels during training. For α, it always start from αstart = 0.5 and gradually decay to αmin in TextWorld experiments. A small-scale parameter search is performed in Difficulty 1 and 4 for αmin. For Difficulty 1, we tried αmin = [0.1, 0.15] and found 0.15 to be the best. For Difficulty 4, we tried αmin =
1https://github.com/xingdi-eric-yuan/GATA-public, MIT License. 2https://github.com/xavierpuigf/virtualhome, MIT license. 3Code available at https://github.com/PKU-RL/MIPO.
[0.1, 0.15, 0.25] and found 0.25 to be the best. For Difficulty 2 and 3, αmin is set to be 0.1 heuristically. In VirtualHome experiments, αmin is set to be 0.1 for both difficulties. Also, for MIPO, a small replay buffer D̃ is used for the language model adaption, we set the size of D̃ to be 12500, which is much smaller than the size of D (500000). For GATA-TrufLL and GATA-CSM, we set the threshold λ = 0.5. For methods using pretrained language models, we use DistilBERT (Sanh et al., 2019) downloaded from Hugging Face. Spacy with en_core_web_sm model is used in all methods, which is consistent with GATA.
For all tables, we report the mean and standard derivations of three different random seeds and bold the result with the best mean performance. For all plots, we plot the curves with the moving average of a window size of 10 and a half standard deviation. The unit of x-axis is 1k episodes. For the generalization results, as MIPO uses a stochastic policy, we evaluate the performance of MIPO 10 times for each generalization task and report the mean.
In terms of the computation budget, all TextWorld experiments are trained in machines rented from online services. All machines are equipped with 12 CPU cores and 1 RTX 3090 GPU. Depending on the algorithm and the difficulty level, each experiment takes around 5 to 45 hours to complete. For VirtualHome experiments, all methods are trained in a local machine with 4 Nvidia GPUs (A100) and 224 Intel CPU Cores. Depending on the algorithm and the difficulty level, each experiment takes around 2 to 5 days to complete.
In terms of the use of AI assistants. ChatGPT is used to polish some sentences in this paper.
A.4 Adaptation to Different Language Model MIPO is a general RL algorithm and can be paired with any pretrained language model, so it should be model-agnostic. We conduct the following experiments to validate this point. In this experiment, we change the pretrained language model to GPT2 for MIPO, GATA-Truff, and GATA-CSM. As we can see from Table 4, MIPO-GPT2 still outperforms other baselines, and it also achieves similar performance with MIPO, which indicates that MIPO is agnostic to the selection of pretrained language model.

heading: Future Work
text: Although we do not observe any significant impact in our experiment, it is still less satisfying and will be addressed in our future work. One another limitation of our work is the value of α, which is currently set heuristically. Although we find it easy to tune in our experiment, a fully automatic method, like the one used in SAC (Haarnoja et al. , 2018b), to set α will be more welcomed and will be included in our future work. To the best of our knowledge, there is no obvious risk to our work