Efficient Reinforcement Learning with Large Language Model Priors
Authors: Xue Yan
Yan Song
Xidong Feng
Mengyue Yang
Haifeng Zhang
Haitham Bou Ammar
Jun Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement Learning; Probabilistic Inference; Language Prior;
Primary area: reinforcement learning
TL;DR: To harness LLMs' rich prior knowledge for efficiently solving complex SDM tasks, we propose treating LLMs as prior action distributions and integrating them into RL frameworks through Bayesian inference methods.

Abstract: In sequential decision-making (SDM) tasks, methods like reinforcement learning (RL) and heuristic search have made notable advances in specific cases. However, they often require extensive exploration and face challenges in generalizing across diverse environments due to their limited grasp of the underlying decision dynamics. In contrast, large language models (LLMs) have recently emerged as powerful general-purpose tools, due to their capacity to maintain vast amounts of domain-specific knowledge. To harness this rich prior knowledge for efficiently solving complex SDM tasks, we propose treating LLMs as prior action distributions and integrating them into RL frameworks through Bayesian inference methods, making use of variational inference and direct posterior sampling. The proposed approaches facilitate the seamless incorporation of fixed LLM priors into both policy-based and value-based RL frameworks. Our experiments show that incorporating LLM-based action priors significantly reduces exploration and optimization complexity, substantially improving sample efficiency compared to traditional RL techniques, e.g., using LLM priors decreases the number of required samples by over 90\% in offline learning scenarios.

Published as a conference paper at ICLR 2025
EFFICIENT REINFORCEMENT LEARNING WITH LARGE LANGUAGE MODEL PRIORS
Xue Yan 12,
Yan Song 3, Xidong Feng" 3, Mengyue Yang Haifeng Zhang? 2 Haitham Bou Ammar 355 Jun Wang!
Institute Automation, Chinese Academy of Science, Beijing, China
School of Artificial Intelligence, University of Chinese Academy of Sciences, China
3 AI Centre, Department of Computer Science, University College London, London, UK University Bristol, UK Huawei Technologies, London, UK
ABSTRACT
In sequential decision-making tasks, methods like reinforcement learning (RL) and heuristic search have made notable advances in specific cases. However, they often require extensive exploration and face challenges in generalizing across diverse environments due to their limited grasp of the underlying decision dynamics, In contrast, large language models (LLMs) have recently emerged as powerful general- purpose tools, due to their capacity to maintain vast amounts of domain-speciio knowledge. To harness this rich prior knowledge for efficiently solving complex sequential decision-makinn tributions and integrating them into RL frameworks through Bayesian inference methods, making of variational inference and direct posterior sampling. The proposed approaches facilitate the seamless incorporation of fixed LLM priors into both policy-basee -baed and value- value-based RL frameworks. Our experiments show that incorporating LLMM-bsse action priors significantly reduces exploration and optimization complexity, substantially improving sample efficiency compared traditional RL techniques, e.g., using LLM priors decreases the number of required samples by over 90% in offline learning scenarios.
tasks, we propose treating L.MSs as prior action dis-
INTRODUCTION
Many real-world tasks, including robotics (Polydoros & Nalpantidis 2017 Brunke et al. 2022 Rana
al. 2023). autonomous driving (Naranjo et al. [2055] [Snng et 2022). human-Al dialogue (Li
et al.
2019: McTear 2022), and human-Al gaming (Granter et al.] 2017: Yan et al. (024), involve
complex sequential decision-making (SDM) challenges. Traditional approaches to SDM, such as optimal control (Garcia al. 1989). heuristic search (Swiechowski al. 2023) and reinforcement learning (RI.) (Mnih) 2013). have seen substantial success. Notably, AlphaGoo (Silver et al.) 2016) and AlphaStar (Vinyals al., 2019), both based on deep reinforcement learning (DRL), have achieved human-level proficiency in the games of Go and StarCraft II, respectively. However, these methods still suffer from high computational complexity, along with poor generalizability and limited applicability across diverse domains (Dulac-Arnold et al. 2015 Cobbe et al. 2019).
Recently, Large Language Models (LLMs) have emerged as effective tools for tackling diverse gereral-purpose tasks, such as in dialogue systems (Brooks et al.) (023). fecissionmmaknn (Zhao) (et al.) 2024a), and mathematical reasoning (Imani et al.) 2023). Their impressive performance across various domains is largely attributed the vast amounts of human knowledge compressed during the pre-training phase on extensive corpora (Tucker& & Tuckute 2023). Inspired by how humans make decisions using existing knowledge while learning from new tasks, LLM-based agents are emerging as promising approach for solving SDM tasks. For instance, researchers leverage human-crafted
"Coreespondence to: (haifeng.zhang (hanradeencctttaaiiaaaaeeeeaaaee @ia.ac.en) ) (jun. wang@es..ul..c.uuk

---

Markoy Decision Process A Markov Decision Process (MDP) can be described as tuple (S, A, P)r,7)- S is the state space, and A the finite action space. We consider the textual action space and state space here. Each state and each action is sentence s, a E y where V is the vocabulary set. S x A + (SS) is the transition kernel, mapping the current state 8; to R is the reward function. 7 is the discount Published as a conference paper at ICLR 2025
prompts to guide LLMs in making decisions directly (Zhang et al. 2023 Wang al. 2023 Ma
ct al. 2023). While this approach relies heavily on the quality of the prompts and the LLMs inherent capabilities, another line of research aims to fine-tune LLMs with RL algorithms (RLFT) (Carta
et al. 2023 Christianos et al. 2023 Tan et al. 2024 /hhou et al. 2024) to enhance their decision- making precision. A more detailed exploration of related work can be found in Section [ However, these approaches are resource-intensiig as they require meticulously crafted human prompts and an expensive fine-tuning process for LLMs for each specific task.
Inspired by these, this work stands in the position of combining the advantages of leveraging rich domain knowledge with the strengths of reinforcement learning (RL), but in a different perspective to use fixed LLMs to enhance traditional RL. The key insight is that while it may be challenging for an L.MM generate optimal plans at every state with limited experience, it can still offer suboptimal action proposals or significantly reduce the size of the action space to a manageable scope, thereby alleviating computational burdens. Building on this idea, we treat the LLM as an action prior distribution and incorporate it into Markov Decision Processes (MDPs) solving from Bayesian inference perspective (Hu et al.). The primary objective is to approximate the posterior action distribution that aligns with task goals through probabilistic inference approaches such variational inference or direct posterior sampling. Building on these approaches, we derive the incorporation of LLMs into both policy-based and value-based RI. frameworks in simple yet logically sound manner. We conduct experiments on major benchmarks like ALFWorld and Overcooked, demonstrating that the sample efficiency of traditional RL algorithms, such as value-based DQN and policy-based PPO, can largely benefit from integrating LLM action priors.
In summary, our main contributions are three-fold:
1. We present unified framework for integrating Large Language Models (LLMs) proba- bilistic priors into Markov decision maasion-making frameworks. (Section
2. We practically implement the framework by leveraging LLMs as the refined action sampler for value-based online RL (Section 3.1 offline RL (Section 3.2) behavior regularizer for policy-based RL (Section 3.3).
3. Extensive experiments on ALFWorld and Overcooked demonstrate that our new framework can significantly boost sample efficiency compared with both pure RL and pure LLM baselines, and also bring in more robust and generalizable value function. (Section Source code is available at hhttps://githu.. com www.yyoouee/e///aaoeaoaaa RL-LLL-LLLPP
2 FORMULATION
We seek to utilize LLMs to solve complex SDM tasks from a Bayesian inference perspective. Although LLM struggles to directly generate optimal plans in sequential decision-making tasks, it can still provide valuable suggestions for possible sub-optimal actions thanks to its rich prior knowledge Hao et al., /Zhan et al. (2024b). Building on this insight, we treat the L.M as a reliable prior distribution over valid actions and analyze its role in solving MDPs from Bayesian inference perspective, using traditional probabilistic tools such as variational inference and probabilistic sampling.
the next state $1+1 following the action 0t. S A factor.
LLM as an action prior It is challenging to directly solve textual decision- tasks due to the lack of task-specific experience, yet powerful LLMs demonstrate the ability to generate reasonable action proposals Hao et al.) Yao et al. (0024). To maximize the potential of LLMs in decision-making
treat the LLM an action sampler, denoted PLLM(a) (a) st). Regarding the implementation details of sampling from LLM priors, it is challenging for an LLM to output an executable action directly. Therefore, prompt the L.LM with the textual state and admissible
2

---

Published as a conference paper at ICLR 2025
actions, then get frre-form output from the LLM (e.g., 7B LLM). This output is subsequently mapped to an executable action through a simple cule-baser projection. formally described as follows: (a),(a)ss) = E. p(a|)) st), where p(a o) is projection from the LLM output action a. More detailed explanations of the L.M prior setting can be found in Appendix 9.2
2.1
VARIATIONAL INFERENCE FOR MARKOV DECISION PROCESS
Similar to the Control as Inference framework proposed by Levine 2018), we introduce binary optimality variable O to indicate the quality of a trajectory
( ag. To, $p.....ppp Spt. where O = 1 represents an optimal/successful trajectory, and O otherwise. Hence, the likelihood, positively related to the reward, can be written as p(O 1/r) x exp (E, 7rr:/a), with as the temperature parameter. Our goal to maximize the optimal marginal distribution. To do that, we first apply variational inference:
pl = 1 ())r) logp(O 1/7)] KL [[(T))(l))(
(1)
We define the prior p(T) the trajectory distribution, which follows the factorization: p(T)
p(so) II. (( 1 1st, 0z). Here, we specify PLLM(als) as the action prior p(als). The vari-
ational distribution (()) follows a similar factorization: ((7) p(5o) II. o(a,/se) (()))) 10,,1111,000 az).
Substituting these factorizations back into Eq. we obtain the step-wise objective as:
arg max
Σ Ei [y're]
aKL (t)
(2)
and we learn a language model policy wy as the variational distribution g(a,|s)) with parameter A. More detailed derivation can be found in Appendix 8.1
However, fine-tuning language models in policy-basee methods can be challenging in practice. To address this, propose alternative approach: direct posterior inference, which allows sample directly from the optimal posterior distribution p(r|O = 1) for sequential decision-making.
2.2 DIRECT POSTERIOR INFERENCE FOR MARKOV DECISION PROCESS
Network: Q(-)
Normalize
a believe
Fixed Bert
Current State
LLM 7B
a Go Go Up
Q(-)
Distribution over prior actions
Adaptor
Next
a goal is...
Up
State
Figure An illustration of the process of approximate sampling from the intractable posterior: q(ass, O 1) p(O l/a, (a)sss(()ss) implemented by reweighting the action prior proposals according to Q-values, which act as likelihood estimates. In our experiments, the Q function adopts BERT encode textual state-action pairs and output a scalar value through adaptor network.
Since the optimal posterior distribution over trajectory T can be factorized into step- wise posterior described as: P(r(O = 1) & p(sp) II, p(ay/st be divided into stages, modeling and inference of step-wise posterior p(a|st, O 1). The modeling stage will point out and decompose the desirable distribution into more solvable probabilistic terms. The inference stage will give the desirable outputs using sampling-bass based techniques (Korbak)
P(s (t-s) Qt-1), the learning procedure
et al. 2022).
Posterior Modeling terms:
We first use Bayes' rule to translate the intricate posterior into accessible
p(a O = 1) p0 1|a, s))((as))) x p(O = 1/a, st)p(a)s)).
(3)
p(O = 1[st)
which is combination of the likelihood p(O la, s,)) and the prior p(a (.)ss).. Here, we use the LLM
the prior action distribution p(als)) PLIM (al st). Recall that the optimality likelihood in MDPs
3

---

Published as a conference paper at ICLR 2025
is given by: p(O 1|a, exp (Ei=r which corresponds to whether the goal state is achieved, similar to the concept of the Q function (8t, a)) in RI. In practice, methods such as MCTS or TDDlearning can estimate the optimality likelihood p(O 1|8;,
Action Inference Having described the posterior as product between LLM prior PLLM (a) and estimate Q function Q" (st, a), we can implement the inference from the posterior distribution p(a|s,, O 1) by following these steps:
Sample actions based on the prior policy PMIM: Obtain i action candidates for the current ... a2.
state st, denoted as: (st) {a1, [01,02...
, where a, (L))(
Choose the next action from the candidates based on the estimated Q-values: Sclect an action a" according to the softmax probability distribution over the estimated Q-values. Specifically, a' ~ softmax(Q" ((t (s),a1) a1)/a, Q a../ hyper-parameter.
ax)/a), where a is
This sampling strategy is illustrated Figure where the action candidates are reweighted based on the Q-values, Below, we establish the connection between variational inference and direct posterior inference.
Proposition 1.
The above action inference strategy stretigy-seeecting action from the action priors reweighted based on Q-values- can be described as following a distribution q. As k - Do, we have:
lim q(als:) = (exx(()ss)expp((? (s, a)/a)// /En,//uullllll)) exp( (s, az)/a)
(4)
The limiting policy corresponds the policy that optimizes the Q-values with a KL regularizer:
lim q(-1st) arg mrg 1ennaxxExtee (=((a)))) (sc, aKL (TK)(((()))))lluun) (((x)s)))luuux((sss)
(5)
Then, the posterior sampling strategy is highly related to the solution of variational inference as shown Eq. 2 Proof. Please see the appendix 8.2
PRACTICAL IMPLEMENTATION
In this section, we introduce the practical implementation of posterior approximation within robust RL frameworks. We describe a policy-based RL method with an LLM prior to realize the variational inference problem. In the meantime, discussed earlier, optimality likelihood estimation critical for direct posterior inference and is positively associated with Q-values in the MDP setting. Thereby, we consider Q-value estimation with an LLLM prior for value-based RI. in both online and offline settings. Finally, we propose three simple yet efficient RI. algorithms: one policy-based RI. algorithm and two Q-learning variants.
We will first introduce the two Q-learning variants, which perform exploration and optimization within narrowed yet reliable LLM prior action space, eliminating the need for fine-tuning the language model-based policy. Then, we will present policy-based PPO variant that incorporates KL loss with respect to the LLM prior. Detailed descriptions of classic Q-learning algorithms, DQN Mnih (2013) and CQL Kumar et al. (2020), can be found in the Appendix
3.1
ONLINE
RL WITH LLM PRIOR
The Q-value estimation is based on DQN Mnih 2013), a widely used value-based RL method, under the online setting. Unlike traditional DQN, which performs exploration and Bellman updates over the full action space A, we propose a variant with an LLM prior, referred to as DQN- This variant limits these processes the prior action space C S A, which is sampled from PLIM This prior action space represents a reduced and more rational subspace where clearly irrelevant actions may be filtered out in advance. Specifically, DQN- -rior alternates between the following steps:
Exploration Roll out new episodes using the posterior inference strategy as described in Sec. 2.2 specifically:
For state 8t, sample the action prior space C² (*((s) {ai, .., ,ax); ) 0: (LLM(a))))
4

---

Published as a conference paper at ICLR 2025
Then sample the action d, softmax (Q" a1)/a........................a.. (st, ax)/a).
Apply the action az to the environment, and obtain the next state $141 ~ P(-)s,,a)):
Add the tuple to the replay buffer: D = (st,, at, 8c+1).
Q-function Update After expanding the replay buffer, we use TD-learnin to update the Q-network. The loss function for the i-th iteration is given as:
()) = E (s,,,))) [10 -s)
(6)
where Yi = (s) r(s, a) + max (5'-'''',a)) a²). Compared to traditional DQN, the DQN-Prior performs exploration and applies the Bellman optimal operator within the LLM prior action space, as highlighted in red in Eq. 6
For most experiments, we use Qwen-1.5 7B Bai al. (2023) generate k prior actions. To represent the Q-network, Bert (Devlin et al. 2019) is used to encode state-action pairs and remains frozen during optimization. Then, introduce an adapter (a 3-layer MLP) on top of Bert's embeddings, as illustrated in Figure ! and optimize the adapter's 's weights by minimizing the temporal difference loss. Therefore, we do not propagate any gradients through the Bert avoid expensive computational overhead.
3.2 OFFLINE VALUE- BASED RL. WITH LLM PRIOR
This study also explores the feasibility of using gradient-free LLM priors for successful SDM based solely on offline datasets, (nell et al. 2023) have explored the use of offline RL for language generation tasks based on the SFT from offline datasets. This approach assumes the entire token vocabulary as the action space, where each action corresponds to single token, resulting in vast action space with space where each action is short phrase consisting of several tokens. For the SDM tasks we address, the action space significantly smaller, with fewer than hundred possible actions. Additionally, by leveraging action prior, can first obtain tidy sub-action space C* thereby narrowing the optimization space.
thousands of possible actions. In our work, consider a discrete action
We apply this action prior within the CQL framework, widely used offline RL approach, and refer to this variant as CQL-Prior. The loss function for the proposed CQL-Prior is given as:
(0) = (s.aa
log Σ ('EC*(a)
+ -
exp(Q"
Qs:
Beea.o' [10 y=)
where y; = r(s, a) + from traditional CQL is that restricts the overestimation of Q-values to the tidy action prior space and applies the Bellman optimal operator within this space, as highlighted in red.
(s', (s), and 3 is a hyper-parameter, The main difference
3.3 POLICY-BASER RL LLM PRIOR
Sample actions
Fixed LLM(7B)
Approximated prior distribution PLAN( [s,)
Current State S,
KL( [s, I1BBuM (s,))
Sample
Action Policy π LLM(<1B)
Next State Reward re
Go
Action distribution from policy T (s
Figure An illustration of the training process of GFlan-Prion Prior, a policy- RL algorithm with an LLM prior. A fixed LLM serves as suboptimal prior, while the action policy learns to approximate posterior action distribution, guided by both environmental rewards and the LLM prior.
5

---

Published as a conference paper at ICLR 2025
We implement variational inference for MDPs with LLMs as action priors using policy-based RL framework. Specifically, we build upon GFlan (Carta et al. 2023). let the Flan- T5 small model (Rae ct al. 2021) (with less than 1B parameters) as the action policy for leveraging pre-stored knowledge to solve complex SDMs and fine-tune via Proximal Policy Optimization (PPO) (Schulman et al. (017). Our variant, called GFlan-Prior, incorporates KL regularizer into the PPO loss to optimize Eq. During implementation, we sample k = 5 examples from LLM prior distribution ( to
compute the approximated KL divergence. Formally, for the i-th iteration, the policy me, is learned by optimizing:
To(arlse) To, (ax))1) (8t)
arg max B(sc,ae)mme CLIP 6
Ae + A - (/oo:=)))))aaL((T////////lllllll
(8)
where A, is the estimated advantage, 17 the hyperparameter that encourages exploration, and PLiM((|ss is the discrete distribution constructed by ic action proposals from the LLM. 0;-1 represents the parameters of the action policy learned in the last iteration. The illustration of the training process of GFlan-Prin can be found in Figure 2
4 EXPERIMENTS
In this section, we demonstrate the effectiveness of incorporating the LLM prior into the RL frame- work, including both policy-based and value-based algorithms in online and offline settings. LLM priors are used to reduce the exploration and optimization space or regulate the action policy's behavior, thereby resulting in efficient RI. framework. Further details on experimental settings and more ablation study results can be found in the Appendix
4.1 ENVIRONMENTS
We consider three environments:
ALFWorld (Shridhar et al. (020) is popular benchmark for examining L.LM-based agents' decision- making ability. This benchmark contains thousands of Textworld (Côté et al. 2019) games with the embodied engine. We focus solely on the text game part, where the action space consists of high-level plans such as "go to a room". The legal actions are finite but large, with the maximum possible admissible action space reaching up to 50, making it challenging to explore from scratch. It contains thousands of tasks, making testing the generalization performance on unseen tasks convenient. We consider two classes of ALFWorld tasks: plF)k)) and ALFWorld(examine) There are no auxiliary rewards, except for a reward of 1.0 for reaching the final goal.
Overcooked We use the partially observed text overcooked game et al. 2024); the observation describes the position of visible items, and the agent should take a sequence of actions to deliver a dish. We consider two overcooked tasks:
crrrcoooed(Tomatoo deliver dish of chopped tomato; Overcooked(Salad) deliver a salad containing chopped tomato and lettuce. The maximum possible action space is 8. Besides the reward of 1 for successfully delivering a dish, the textual Overcooked environment from (Tan et al. 2024) also provides dense reward signals. The reward shaping is as follows: 0.2 for correctly chopping an ingredient, terminal reward for successfully delivering the correct dish, 1 for delivering any incorrect item, and -0. -0.001 for each time step.
Frozen Lake is a grid world game; the agent should move the goal position while avoiding falling into holes. There are four admissible actions: up, down, right, and left. The agent will receive
reward of 1 for reaching the final goal and a penalty of -1 for falling into holes. More details of state and action representations can be found in Appendix 9.1
4.2 BASELINES
We first test the ability of LLM Prior to solve SDMs in a zero-shot manner, without deliberately designed prompts. Next, we will introduce the value-based and policy-based RL algorithms:
Salued-Based RL: For the online setting, we use DQN and our DQN-Prior, which explores and updates the Q-function within the LLM prior space. For the offline setting, we compare the CQL, CQL-prior, and Behavioral Cloning(BC). Similarly, compared to CQL, CQL-Prior regulates and updates Q-values in the narrowed prior action space. We set l 5 for DQN-Prie and CQL-Prion The BC utilizes the Flan- small (Rae et al. 202T) as the action policy, and learns the action policy by optimizing: arg min, E. (..)D log 77

---

Published as a conference paper at ICLR 2025
DON-Prior(ours
GFlan
ALFWorld (Pick)
ALFoddrid (Examine)
Overzooked (Tormato)
1.00
0.8
0.75
Roward
Reward
0.50
0.25
0.2
0.00
0.25
0
200
400
600
800
1000
Episodes
200
400
600
800
1000
50
100
150
Episodes
Eplsodes
Overcooked (Salad)
1.0
0.8
Roward
0.5
Reward
0.6
0.0
0.2
0.5
0.0
0
50
100
150
Episodes
Episodes
Figure 3: Results of comparison with online baselines. We plot the mean and standard error of the cumulative reward. Periodic evaluation results are plotted for online baselines.
Policy-Based RL: GFlan (Carta et al.) 2023 lets the LLM Flan-T5 small the action policy and finetunes PPO (Schulman et al. 2017). GFlan-Prior adds a KL constraint between optimized action policy and LLM prior action on the basis of GFlan, shown in Eq. 26
LLM Setting Due to the Qwen- 1.5 series (Bai et al. 2023) containing a range of LLMs with different scales, use Qwen-1.5 for our main experiments. Specifically, we use Qwen-1.5 7B (Bai ct al. 2022)) as the backbone of the LLM prior for most environments, except for the more complex (vercooked(Sala)) task, where Qwen-1.5 14B is used.
4.3 RESULTS ANALYSIS
4.3.1
INCORPORATING LLM PRIORS INTO VALUE-BASED RL
The effectiveness of LLM Prior on online Q-learning As shown in Figure 3 our DQN-Prior outperforms traditional online RL baselines such DQN and GFlan by exploring and optimizing within a narrowed, more valuable prior action space. In complex ALFWorld tasks. DQN and GFlan struggle bootstrap due the large original action space, which can include up 50 possible actions and the lack of dense reward incentives. In contrast, our DQN-Prior leverages the L.M's domain knowledge to reduce the exploration space to just five suboptimal actions, significantly improving sample efficiency. Additionally, in the Overcooked (Salad) and Frozen Lake tasks, the LLM prior cannot directly solve these tasks by sampling one action per state. However, DQN-Prior still achieves success and demonstrates greater sample efficiency than traditional RL methods, despite the small action space (only 4 or 8 actions). This shows that the LLM prior can provide suboptimal action proposal space, and conducting RL within this space significantly enhances sample efficiency.
Table 1: Results of offline algorithms. We consider two tasks, ALFWorld (Pick) and Overcooked (Salad), referred to Pick and Salad for simplicity. To further compare the sample efficiency of the basclines, we use offline datascts varying sizes for Overcooked (Salad), denoted as Salad (N), each containing approximately N (s, a.) tuples.
Baseline
Pick Salad(1000) Salad(4000)
Salad(8000)
Salad(12000)
Salad(24000)
DQN
0.27
0.14
0.31
0.32
0.32
0.32
CQL
0.03
0.32
0.32
0.32
0.32
1.33
DON-Prior
0.49
0.98
0.81
CQL-Prior
0.80
1.01
1.19
1.31
BC
0.31
0.57
0.91
0.78
0.53
0.74
7

---

Published as a conference paper at ICLR 2025
Table 2: Results the generalization performance of posterior sampling on AlFWorld(Pick)) We use Qwen- 1.5 7B with & = 5 to generate LLM prior action proposals for training the Q function via DQN-Prior and CQL-Prior. This LLM prior configuration is highlighted by (*).
Online (DQN-Prior)
Seen Tasks
Unseen Tasks
k
k =
5
k = 10
k = 15
k =
1
k =
5
10
A
15
Qwen- 4B
0.35
0.46
0.35
0.08
0.25
0.29
Qwen-1.5 7B
0.19
0.62
0.77
0.65
0.04
0.38
0.42
0.42
LLaMa-3 8B
0.04
0.62
0.73
0.69
0.13
0.42
0.38
0.38
Qwen-1.5 14B
0.19
0.65
0.77
0.65
0.16
0.54
0.54
0.38
Qwen-1.5 32B
0.35
0.77
0.81
0.73
0.33
0.5
0.5
0.46
Offline (CQL-Prior)
Seen Tasks
Unseen Tasks
k =
k
k
10
k = 15
k ==
k =
5
10
k 15
Qwen- 4B
0
0.27
0.38
0.62
0.0
0.17
0.17
0.38
Qwen- 7B
0.19
0.80*
0.81
0.77
0.04
0.46
0.50
0.50
LLaMa-3 8B
0.04
0.85
0.88
0.81
0.13
0.54
0.50
0.54
Qwen- 14B
0.19
0.73
0.92
0.73
0.16
0.54
0.58
0.54
Qwen-1.5 32B
0.35
0.81
0.85
0.69
0.33
0.63
0.42
0.50
The effectiveness of LLM Prior offline Q-learning Table illustrates the performance of baselines trained on offline datasets. A more detailed description of the offline datasets can be found in the Appendix. Table T shows that CQL-Prior outperforms all other baselines on the ALFWorld (Pick) and Overcooked (Salad) datasets. Although DQN- -oiorrior lacks constraints on Q-values, it still demonstrates superior performance compared DQN and CQL, which operate on the full action space. These results indicate that in offline RI,, avoiding overestimation of Q-values and performing Bellman updates only within the suboptimal prior action space reduces optimization complexity, thereby improving sample efficiency.
Additionally, we provide further evidence to support this conclusion by testing the sample require- ments for CQL and CQL-Prior. We find that traditional CQL, with fewer than 12000 examples, fails to learn how to achieve the final goal, only managing to reach the subgoal of chopping the tomato and lettuce (with reward of 0.4). In contrast, CQL-Prior, with just around 1000 examples, successfully learns how to achieve the final goal (with a reward 1.0) with high probability. Therefore, our CQL-Prior reduces the number of required samples by at least 90% compared to CQL this task.
The generalization ability across unseen tasks and LLMs Table demonstrates the generalization ability of our posterior sampling framework by combining the Q-function, trained with the LLM prior setting of Qwen-1.5 7B and & = 5, with other LLMs and varying action proposal numbers (k) during inference. Even though the Q-network is trained with Qwen-1.5 7B and A 5, it can still effectively guide the decision-making process, i.e., action selection, for unseen LLMs of different scales and architectures. It is worth noting that, although Qwen-1.5 4B was originally unable to solve ALFWorld (Pick), it becomes capable of tackling such complex tasks by leveraging the Q-function trained with more powerful LLMs. This value-based RL approach, which utilizes LLM prior, presents a promising opportunity to train an information-rich Q-network using larger LLMS while deploying only the Q-network and smaller LLMs on the client side. This method enhances reasoning speed and reduces deployment complexity. Furthermore, larger LLMs generally perform better during inference, as their action proposals are of higher quality. k = indicates the performance of the L.M on its own, without involving the Q-function. The setting of l = 5 was chosen during training for querying efficiency, but it is not necessarily the optimal choice. We observed better test performance with & = 10 seen tasks, indicating that our method can generalize to larger action space to some extent. In most cases, inference performance increases and then decreases as k increases. This may occur because larger action space is more likely to include the optimal action, but excessively large action space can introduce unseen (s, a) pairs during training, leading to of distrribut (OOD) problems in both online and offline settings.
8

---

Published as a conference paper at ICLR 2025
ALFWorld (Pick)
Overcooked (Tomato)
Overcooked (Salad)
0.8
1.0
DQN Prior
1.0
0.8
GFian
Roward
LLM-Prior
0.6
Rewerd
0.4
GFlan-Prior
0.4
0.2
0.2
0.0
-0.5
0.0
200
400
800
800
1000
150
Episades
Episodes
Episodes
(a) Results of policy- based RL with LLM prior
ALFoodl (Pick)
Overcooked (Tomato)
Overcooked (Tomato)
0.8
DQN-Priar
a= 0.0
1.0
0.6
GFlen
0.8
a 0.01
0.8
Roward
GFlan-Pror((e=))
a= 0.1
0.4
GFllan-Prior(k--)
0.4
1.0
0.2
0.2
0.0
0.0
0.2
200
400
600
800
1000
0
50
100
150
50
100
150
Episodes
Episodes
Episodes
(b) GFlan-Prior
(c) a for GFlan-Prior
(d) DQN- -rior
Figure 4: (a) Comparison of policy-based RI algorithms, plot DQN-Prior and L.M--rrior for reference. (b) The ablation study on the number of L.M action proposals k used to approximate the KL divergence. (c) The ablation of the KL coefficient of Gllan-Prior. (d) The ablation of the softmax temperature over Q values of the DQN-Prior
Regarding generalization to unseen tasks, multiple LLM action proposals with A: 1, guided by the O-function, generally outperform the LLM's ' zero-shot performance with i = 1. This demonstrates the generalization capability of our posterior sampling framework, enabled by the powerful LLM as an action prior and the Q-function's ability to compress interactive experiences.
4.3.2 INCORPORATING LLM PRIORS INTO POLICY-BASED RL
The Figure 4(a)
show the results policy-based RL algorithms. As shown in Figure ((a) GFlan-Prior significantly outperforms GFlan and achieves performance comparable DQN- Pri-rrior. Unlike GFlan, GFlan-Priae Pritann-Prioooa leverages the prior knowledge embedded LLMs by aligning with the LLM prior distribution through the incorporation of KL constraint, shown in Equation 26 In our main results, we sample &o = action proposals from the LLM prior to approximate the KL divergence. An ablation study & can be found Figure (bb) and Figure in the Appendix. Similar to the results Table with k = 5, we are more likely to sample the optimal action than with k 1, thereby better supervising the learning of the action policy.
4(b)
4(c)
and
The ablation study on the coefficient of the KL divergence is shown in Figure A large will cause the learned policy æ to closely follow the LLM prior proposals, while small a encourages exploration. In DQN-Prior, the hyperparameter regulates exploration by acting as the temperature in the softmax over Q, generating Boltzmann distribution based on Q, as investigated Figure Kd) GFlan-Prior appears to be more sensitive to the hyperparameter controlling exploration than DQN-Prior.
RELATED WORK
This study explores the fusion of Bayesian inference and L.LMs for decision-making tasks. Herein, we present related works on LL.M-Agent for decision- makinn tasks and how previous research pursues this combination. More discussions on the use of L.LMs for value-based inference and reward function generation are provided in Appendix I0
LLM-driven Agent The idea of LLM-as-Agemt has gone viral since the release of powerful large language models. Among all the works, some apply direct prompt engineering techniques

---

Published as a conference paper at ICLR 2025
using human--esignn prompt and retrieval to harness LLMs for complex interactive tasks, such as ProAgent by (Zhang et al. 2023). Voyager by (Wang et al. 2023) and generative agent by (Park) et al. 2023). Many researchers have also tried searching-based methods. (Yao et al.) 2023) propose Sete-of-TToughle (TOT) using BFS/DFS algorithm to search accurate decision sequences. (Hao et al.)
suggests using Monte-Carlo Tree Search to perform more comprehensive search. These methods often fall into the framework of Proposer-Verifier Verifier (Snell et al. (2024) where the language proposes a number of possible sequences and the verifier picks desirable candidates. There are also few works that relate LLM to conventional reinforcement learning settings and investigate how the superiority decision-making For instance, (Brooks et al. 2024) uses LLM world model and performs policy iteration by in-context learning. GFlan proposed by (Carta et al. 2023) uses a learnable LLM as probability likelihood estimator for all possible actions and incorporates it into actor-critic learning frameworks where value function generates assessment and fine-tuning the LLM correspondingly. (Yan et al. 2023) also resort to an actor-critic learning paradigm but instead of fine-tuning the language model, they train a simple classifier that selects valuable outputs from a fixed language model (e.g. GPT3.5). (Zhang et al.) 2024b) use LLM as a behavior regularize and add it to the value estimate. (Wen et al. 2024h leverages entropy-regulariie reinforcement learning to train LLM agents for verbal sequential decision-making tasks. Our work also falls into the category of LLM-driven RL agent. However, different from others, we see large language model as an excellent action proposer and perform Q-learning directly on the reformed action sequence space, thereby achieving better learning efficiency with minimum modification.
in language space helps RL agents' capabiliity
LLM in Bayesian Sense The probabilistic nature of transformer-basee language models makes them well-suited for Bayesian inference framework. As highlighted by (Korbak et al.. 2022) the Bayesian approach offers a unified perspective on both modeling and inference. in-context learning (ICL), some researchers argue for the Bayesian properties of ICI. (fJingg 2023 Wang et al. 2024b Ye ct al. 2024)). Extending beyond ICL, (Yang & Klein 202T) trains a classifier to approximate the likelihood function, framing conditional text generation as a Bayesian inference problem, with the
LLM serving as the prior. (Hu et al.) interprets Chain-of- Chain-of-Thought Thought (CoT) reasoning in natural language processing (NLP) tasks, such text infilling and sequence continuation, as probabilistic inference problems, using GelowNet to fine-tune the model. In this context, the language model functions as a unified probabilistic generative process, capable of representing both prior and likelihood distributions. Similarly. (Zhao et al. (2024b) adopts Bayesian perspective, utilizing Sequential Monte Carlo (SMC) methods to generate undesired outcomes. (Gallego) (2024) frames LLM inference as sampling from a posterior distribution focused on high-reward regions, applying this approach to self-improvement tasks. Inspired by these works, we also formulate our agent framework within a Bayesian setting, exploring the use of LLMs as priors.
CONCLUSION
In this work, we aim to harness the rich domain knowledge and strong reasoning abilities of LLMs to tackle complex pocision-making tasks while avoiding the costly processes of prompt design and fine-tuning. Given their extensive prior knowledge but limited task-specific experience, we do not rely on LLMs to directly generate optimal plans. Instead, we simplify their role to producing reliable action proposals. To this end, we treat the powerful LLM as an action prior distribution and, from Bayesian inference perspective, examine how integrate LLM priors into solving MDPs. We employ variational inference and direct posterior sampling to achieve this, ultimately proposing both policy- based and value-based RL frameworks incorporating LLM priors. These priors enhance traditional RL frameworks by narrowing the exploration space or guiding the action policy through suboptimal --gnerrre proposals. Experimental results demonstrate a substantial improvement in sample efficiency when incorporating fixed L.M priors into RI. frameworks in both online and offline settings. This work focuses specifically on text-based games with finite action spaces. In future research, we plan to explore scenarios with free-form and infinite action spaces and investigate deliberate reasoning approaches to further refine the quality of --M-geeeetteed action proposals.
ACKNOWLEDGEMENTS
Co-author Haifeng Zhang thanks IPT Project 2024002.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023.
Ethan Brooks, Logan A Walls, Richard Lewis, and Satinder Singh. Large language models can implement policy iteration. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Ethan Brooks, Logan Walls, Richard I. Lewis, and Satinder Singh. Large language models can implement policy iteration. Advances in Neural Information Processing Systems, 36, 2024.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoe V Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXiv:2407.2 1787, 2024.
Lukas Brunke, Melissa Greeff, Adam W Hall, Zhaocong Yuan, Siqi Zhou, Jacopo Panerati, and Angela Schoellig. Safe learning in robotics: From learning-based control to safe reinforcement learning, Annual Review of Control, Robotics, and Autonomous Systems, 5(1):41 1-4444 2022.
Thomas Carta, Clément Romac, Thomas Wolf, Sylvain Lamprier, Olivier Sigaud, and Pierre- Yves Oudeyer. Grounding large language models in interactive enviromments with online reinforcement learning. International Conference on Machine Learning, pp. 3676-3713, PMLR, 2023.
Filippos Christianos, Georgios Papoudakis, Matthieu Zimmer, Thomas Coste, Zhihao Wu, Jingxuan Chen, Khyati Khandelwal, James Doran, Xidong Feng, Jiacheng Liu, Zheng Xiong, Yicheng Luo, Jianye Hao, Kun Shao, Haitham Bou-Ammar, and Jun Wang. Pangu-agent: A fine-tunable generalist agent with structured reasoning. 2023.
Karl Cobbe, Oleg Klimov, Chris Hesse, Tachoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In International conference on machine learning, pp. 1282-1289. PMLR, 2019.
Marc-Alexandre Côté, Akos Kádár, Xingdi Yuan, Ben Kybartas, Tavian Barnes, Emery Fine, James Moore, Matthew Hausknecht, Layla Asri, Mahmoud Adada, et al. Textworld: A learning environment for text-based games. In Computer Games: 7th Workshop, CGW 2018, Held in Con- junction with the 27th International Conference on Artificial Intelligence, IJCAI 2018, Stockholm, Sweden, July 13, 2018, Revised Selected Papers 7, pp. 41-75. Springer, 2019.
Jacob Devlin, Ming-Wor Wei Chang, Kenton Lec, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pp. 4171-4186, 2019.
Gabriel Dulac-Arnold, Richard Evans, Hado van Hasselt, Peter Sunchag, Timothy Lillicrap, Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, and Ben Coppin. Deep reinforcement learning in large discrete action spaces. arXiv preprint arXiv: 152.0077, 2015.
Fares Fourati, Vaneet Aggarwal, and Mohamed-Slim Alouini. Stochastic q-learning for large discrete action spaces, arXiv preprint arXiv:2405. 10310, 2024.
Victor Gallego. Distilled self-critique of Ilms with synthetic data: a bayesian perspective. In The Second Tiny Papers Track at ICLR 2024, 2024.
Carlos Garcia, David M Prett, and Manfred Morari. Model predictive control: Theory and practice- survey. Automatica, 25(3):335-344 1989.
Scott R Granter, Andrew H Beck, and David J Papke Jr. Alphago, deep learning, and the future of the human microscopist. Archives of pathology & laboratory medicine, 141(5):619-62) 6211,,11151 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference Machine Learning (ICMI.) pp. 1861-1870. PMLR, 2018
11

---

Published as a conference paper at ICLR 2025
Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu. Reasoning with language model is planning with world model. In The 2023 Conference on Empirical Methods in Natural Language Processing.
Edward J Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, and Nikolay Malkin. Amortizing intractable inference in large language models. In The Twelfth International Conference on Learning Representations.
Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang. Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2021.
Shima Imani, Liang Du, and Harsh Shrivastava. Mathprompter: Mathematical reasoning using large language models. arXiv preprint wrXivv2303.05396 2023.
Hui Jiang. A latent space theory for emergent abilities in large language models. arXiv preprint arXiv:2304.0990 099600 2023.
Martin Klissarov, Pierluca D'Oro, Shagun Sodhani, Roberta Raileanu, Pierre-Lue Bacon, Pascal Vincent, Amy Zhang, and Mikael Henaff. Motif: Intrinsic motivation from artificial intelligence feedback. arXiv preprint arXiv: 10.00166, 2023.
Tomasz Korbak, Ethan Perez, and Christopher L Buckley, RI with penalties is better viewed as bayesian inference. arXiv preprint urXiv:2205.11277 2022.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances Neural Information Processing Systems, 33:1 179-1191, 2020.
Minae Kwon, Sang Michael Xie, Kalesha Bullard, and Dorsa Sadigh. Reward design with language models. In The Eleventh International Conference on Learning Representations, 2023.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: (rXXiv1805500900 2018
Kenneth Li, Samy Jelassi, Hugh Zhang, Sham Kakade, Martin Wattenberg, and David Brandfonbrener. Q-probe A lightweight approach to reward maximization for language models. arXiv preprint arXiv:2402. 14688, 2024.
Ziming Li, Julia Kiseleva, and Maarten De Rijke. Dialogue generation: From imitation learning to inverse reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 6722-6729, 2019.
Hunter Lightman. Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint arXiv:2305.20055 2023.
Weiyu Ma, Qirui Mi, Xue Yan, Yuqiao Wu, Runji Lin, Haifeng Zhang, and Jun Wang. Large language models play stareraft ii: Benchmarks and chain of summarization approach. arXiv preprint
arXiv:2312.11865. 2023.
Yecheng Jason Ma. William Liang, Guanzhi Wang, De-An Huang, Osbert Bastani, Dinesh Jayaraman, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Eureka: Human-leved reward design via coding large language models. In The Twelfth International Conference Learning Representations, 2024.
Nat McAleese, Rai Michael Pokorny, Juan Felipe Ceron Uribe, Evgenia Nitishinskaya, Maja Trebacz, and Jan Leike. Llm critics help catch I1m bugs. arXiv preprint drXiv:2407.00215, 2024.
Michael McTear. Conversational ai: Dialogue systems, conversational agents, and chatbots. Springer Nature, 2022.
Volodymyr Mnih. Playing atari with deep reinforcement learning. arXiv preprint urXiv:1312.500 2013.
12

---

Published as a conference paper at ICLR 2025
José Eugenio Naranjo, Carlos González, Ricardo García, Teresa de Pedro, and Rodolfo E Haber. Power-steering control architecture for automatic driving. leee transactions on intelligent trans- portation systems, 6(4):406 415, 2005.
Joon Sung Park, Joseph D'Brien, Carrie Jun Cai, Meredith Ringel Morris, Percy Liang, and Michael Bernstein. Generative agents: Interactive simulacra of human behavior. In Proceedings of the 36th Annual ACM Symposium on User Interface Software and Technology. pp. 1-22, 2023.
Athanasios S Polydoros and Lazaros Nalpantidis. Survey of model-based reinforcement learning: Applications on robotics. Journal of Intelligent Robotic Systems, 86(2):153 173, 2017.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11444 2021.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. arXiv preprint arXiv:2305. 18t.:2300...82990, 2023.
Krishan Rana, Ming Xu, Brendan Tidd, Michael Milford, and Niko Sünderhauf. Residual skill policies: Learning an adaptable skill-based action space for reinforcement learning for robotics. In Conference Robot Learning, pp. 2095-2104. PMLR, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1007.0/7777.634777 2017.
Mohit Shridhar, Xingdi Yuan, Alea-ndex Côté, Yonatan Bisk, Adam Trischler, and Matthew Hausknecht. Alfworld: Aligning text and embodied environments for interactive learning. arXiv preprint wwxx::200000076660 2020.
David Silver, Aja Huang, Christopher J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schritwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, llya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529:484-503, 2016. URL http / www. , www.rattu.....o.oa htt.://www....oooo../../////////
html
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can more effective than scaling model parameters. arXiv preprint 2xXiv:2408.03314 2024.
Charlie Victor Snell, Ilya Kostrikov, Yi Su, Sherry Yang, and Sergey Levine. Offline rl for natural language generation with implicit language q learning. In The Eleventh International Conference on Learning Representations, 2023.
Qingyuan Song, Weiping Fu, Wen Wang. Yuan Sun, Denggui Wang, and Jíncao Zhou. Quantum decision making in automatic driving. Scientific reports, (2(1):11042, 2022.
Maciej Swiechowski, Konrad Godlewski, Bartosz Sawicki, and Jacek Mańdziuk. Monte carlo tree search: A review of recent modifications and applications. Artificial Intelligence Review, 56(3): 2497-2562, 2023.
Weihao Tan, Wentao Zhang, Shanqi Liu, Longtao Zheng, Xinrun Wang, and Bo An. True knowl- edge comes from practice: Aligning large language models with embodied environments via reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024.
Mycal Tucker and Greta Tuckute. Increasing brain-llm alignment via information-theorriic compres- sion. In UniReps: the First Workshop on Unifying Representations in Neural Models, 2023.
Oriol Vinyals. Igor Babuschkin, Wojciech M Czamecki, Michaël Mathieu, Audrey Dudzik, Junyoung Chung, Arthur J Tashman, David Horgan, Peter Kroiss, Max Powell, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)3300-35. 2019.
13

---

Published as a conference paper at ICLR 2025
Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint ar arXiv:2305..62911 2023.
Peiyi Wang, Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce Ilms step-by-step without human annotations. In Proceedings of the 62nd Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 9426-9439, 2024a.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are latent variable models: Explaining and finding good demonstrations for in-context learning. Advances in Neural Information Processing Systems, 36, 2024b.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdh- ery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11177 2022.
Muning Wen, Cheng Deng, Jun Wang, Weinan Zhang, and Ying Wen. Entropy- token-level policy optimization for large language models. arXiv preprint arXiv:2402.06700, 2024.
Xue Yan, Yan Song, Xinyu Cui, Filippos Christianos, Haifeng Zhang. David Henry Mguni, and Jun Wang. Ask more, know better: Reinforce-leamed prompt questions for decision making with large language models. arXiv preprint 0rXiv:2310.88122 2023.
Xue Yan, Jiaxian Guo, Xingzhou Lou, Jun Wang, Haifeng Zhang, and Yali Du An efficient end- to-end training approach for zero-shot human-ai coordination. Advances in Neural Information Processing Systems, 36, 2024.
Kevin Yang and Dan Klein. Fudge: Controlled text generation with future discriminators. arXiv preprint wxXiv:2104.05218, 2021.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. arXiv preprint arXiv:2305. 10601, 2023.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
Naimeng Ye, Hanming Yang, Andrew Siah, and Hongseok Namkoong. Pre-training and in-context learning is bayesian inference la de finetti. arXiv preprint arXiv:2408. 03307, 2024.
Wenhao Yu, Nimrod Gileadi, Chuyuan Fu, Sean Kirmani, Kuang-Hus Huei Lee, Montserrat Gonzalez Arenas, Hao-Tien Lewis Chiang, Tom Frez, Leonard Hasenclever, Jan Humplik, et al. Language to rewards for robotic skill synthesis. In Conference Robot Learning, pp. 374-4044 2023.
Ceyao Zhang, Kaijie Yang, Siyi Hu, Zihao Wang, Guanghe Li, Yihang Sun, Cheng Zhang, Zhaowei Zhang, Anji Liu, Song-Chun Zhu, al. Proagent: Building proactive cooperative ai with large language models. arXiv preprint arXiv:2308. 1339, 2023.
Lunjun Zhang, Arian Hosseini, Hritik Bansal, Mehran Kazemi, Aviral Kumar, and Rishabh Agarwal. Generative verifiers: Reward modeling as next-token prediction. arXiv preprint arXiv:2408.15244 15240, 2024a.
Shenao Zhang, Sirui Zheng, Shuqi Ke, Zhihan Liu, Wanxin Jin, Jianbo Yuan, Yingxiang Yang. Hongxia Yang, and Zhaoran Wang. How can Ilm guide rl? value- based approach. arXiv preprint arXiv:2402.16188 16181, 2024b.
Andrew Zhao, Daniel Huang, Quentin Xu, Matthieu Lin, Yong-Jin Jin Liu, and Gao Huang. Expel: Llm agents are experiential learners. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 19632-19641 2024a.
14

---

Published as a conference paper at ICLR 2025
Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Baker Grosse. Probabilistic inference in language models via twisted sequential monte carlo. In International Conference on Machine Learning, pp. 60704-60744 PMLR, 2024b.
Runlong Zhou, Simon Du, and Beibin Li. Reflect-rl: Two-player online rl fine-tuning for Ims. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1015 2024.
15

---

Published as a conference paper at ICLR 2025
7 DESCRIPTION OF CLASSIC Q-LEARNING ALGORITHM
We introduce the classic Q-lcarning algorithm both offline and online settings.
Q-learning Q-learning is classic value-based RL algorithm, which iteratively applies the Bellman optimal operator to train the Q-function, formally written as B ({{ = r(s,a) + ((.la)(())s.a) [max'x Q(s', " d')] After the optimal Q-function Q" learned, then the optimal action pol- icy respect to the Q* is given as: π (-Is) arg max, Q* a). DQN (Mnih] 2013) is well-known Q-learning algorithm, which can process various classes of the state information such as image and language, which learns deep neural network Q (s. a) = (s, a) approximate Q" (s, a) and uses TD-leaming to iteratively update the Q network, the loss function for the i-the iteration is given as:
(0) ()) [[Q% (s. 3
where 'D is the reply buffer, Vi Q (s, a) r(s, a) + y max.'ed (s' a'). 0;-1 is the learned parameters of the last iteration and is frozen during the gradient descent for optimizing the loss function Li.
Conservative Q-learning CQL. (Kumar et al. 2020) framework is designed for offline reinforce- ment learning, where only offline data is available, but further online exploration is accessible. Different from traditional Q-learning, CQL introduces an additional regularization term on top of the standard DQN framework actor-critic franework(such SAC (Haarnoja et al. 2018)), ensuring that the learned Q-function lower-bounds the actual value. The lower-bounded Q-values help mitigate the overestimation problem commonly encountered with out-of diitr--disttiuttion (OOD) state-action pairs. The CQL loss on the top of value-based DQN is given as:
(0) = (a,)~D log Σ
exp(Q"(ss,d)) - Q (s,a) + B ( - v1)3]
a'EA
(10)
where y: = r(s,a) + max Q (s',a') and B is hyper-parameter.
PROOFS
8.1 PROBABILISTIC INFERENCE DERIVATION
In this section, we derive the detailed formulation for probabilistic inference for using LLM as a prior. This includes the policy-based method where we learn parametrized policy function and the value-based method where we directly perform posterior inference.
Variational Inference Assume that we have the optimality variable O indicating the quality of a trajectory т = {80, 40, 81,..., 8n) .P(T) p(so) Il p(a, (s))p(((pp))) (p()))(((111 8e-1/81. 0z). The likelihood function with temperature parameter can be written as:
p( 1/7) exp Erinja
(11)
Our goal is to maximize the optimal marginal distribution formulated as follows:
p(O = 1) 1(0 = =1,T)dr
(12)
f 9(()) pl0 ((())((0 1|T)p(T) dr
g(T)
where ((7) parametrized variational distribution. Then, our variational inference objective can be written as:
1o = 1) / (() p(O = 1/7) P(T)
dr
((T)
(13)
= Egtt)) ((O = 1/7)] - KL [g(7))p(T)]
ELBO
16

---

Published as a conference paper at ICLR 2025
In the context of language models, we define prior p(r) as trajectory distribution generated through fixed LLM denoted as:
(L = p(so) II (a (a)ss)/((((ss11
(14)
where PI.IM (ar is the probability likelihood of the generated output at given input text S;. The variational distribution ((T) can also be decomposed in the same way as:
<(T) = p(00) II no (a (s)p(8+++1
(15)
where the parametrized policy function Te can be learnable language model with parameter Substituting these factorizations back to ELBO we have the step-wise objective function written as:
- C = Σ Eso [n'r:]
aKL (ro(acss))lluzzaa(aaaa
(16)
π
arg min L
(17)
Direct Posterior Inference Instead of learning parametrized policy which might involve fine-
tuning a language model, we can also try directly sampling from the posterior distribution formulated as:
p(T - p(O 1T7)
p(0 1)
p(O 1/T)p(T)
(18)
II p(O, 1/8t, Oe
According to centrol-as-Inferemn framework proposed by Levine (2018), we can formulate posterior inference as Soft Q-Learning where Q-values are updated using modified soft Bellman equation:
in which define P(T) in Equation
14
Q" (st az) r(s), + +yyB 1-Ber1pp(((ss (V(+++))
V(s)) = log / exp((((s,
(19)
and the policy can defined as a softmax function over the Q-values: <[als) exp(Q" (s; /)/a)/ Jar (x,(("(ssa')) ) ensuring that actions with higher Q-values are more likely be chosen. In our paper, however, due the enormously large space of (textual) state-action pair, we also apply sampling approximation methods. Specifically, we leverage similar ideas from Fourati et al. (2024) which sample a random subset of actions from the complete action space and seek the optimal within this subset. In our case, we rely on fixed LLM to provide such action subspace and can further narrow down the space through repeated sampling (Brown al. (2024)):
LEcPLEM(((s) exp (Q*Ts.a)/al T(lss) Jo. (06p))M((10) exp (Q=(s a/) da'
(20)
1 exp(Q=(s,,')/a) lee(a N N expp("(ss,'')))
PL.MM(als)
This means we restrict the Bellman backup to specific action subspace that can potentially provide a refined set of actions based on context, largely reducing computational complexity and focusing on more relevant actions.
8.2 PROOF OF PROPOSITION
1
Proposition 1. The above action inference strategy-selectin an action from the action priors
reweighted based on Q-values-cun can be described following distribution q. As k - co, we have:
lim q(als:) = Puls{aass))xx((C (s, a)/o)/ (aao))//a,muuue (s,
(4)
100
17

---

Published as a conference paper at ICLR 2025
The limiting policy corresponds to the policy that optimizes the Q-values with KL regularizer:
lim q(-1se) arg max (x(a)s.) (st;
(st; -KKL (x((-1s)) ((r(((s))lpuuu(((ss)
(5)
1+000
Then, the posterior sampling strategy is highly related the solution of variational inference as shown in Eq. Proof. Please see the appendix
Proof. The proof mainly follows the (Li et al., 2024). Denote the above sampling strategy indeed follows distribution and for each action a, the probability of is sampled from the distribution given by:
T(a; a) exp(Q" a;)/a) h
(a))s)) = ER (ar.... az dp3rpux(-11)
Σ exp(X (5L, az)/a))
E I(a; a)
E (ar.""agg))nn
exp(Q
(21)
j exp(Q" (s2, a,)
E Mia =
exp(Q"
Eド= exp exp(pp (st, ag)/a)
According to the Law of Large Numbers, when h: co, we have:
Σ Ila; -
lim q(a(s)) = lim (a). (al.:raa)/muu(i
exp(Q2 a)/a)
80000
x 2= exp((x az)/a)
PLLM(a|sE)
lim E aa an (ap
E Bapppin((as) (x, a;)/a)
k-000
exp(Q aa)
= =PIM((((s))
(as-suu((1a)) (xp(Q" (s)
(22)
Following the proof process from the Appendix A.1. in Rafailov al. (2023), have:
exp(O"
arg max 18 (a)=)) 1Q -)L (T((ssa)ll (als(a)s:)
(ap (s,p(Q((,,,)))
(23)
Additionally, as shown in Sec. maximising:
2 the variational inference approach learns the optimal policy æ by
arg max
Σ E [n're] (ililaaa(((ass))lllllllll...
(24)
Introducing the occupancy measure p. p(s) = L = s/a ~ T)}. we have the following form objective respected to the Q-function
arg max E p(s) ( a)] - aKL(((a)s:)) (a
(25)
Combining Equations 23 and 25 we find that the variational inference approach and the direct
posterior sampling yield similar solutions.
ADDITIONAL EXPERIMENT DETAILS
9.1 ENVIRONMENTS
We consider three environments: ALFWorld, Overcooked and Frozen Lake. ALFWorld We consider two classes of ALFWorld tasks: ALFWorld(pick) and ALxaonnd(examine) For ALFWorld (Pick), we evaluate the online training baselines on 28 tasks, such as "put the cellphone on the armchair," " and test the generalization ability 26 unseen tasks. For Forld (Examine),
18

---

Published as a conference paper at ICLR 2025
AlFcorld(Picc)
Xxxmine)()aamine)
Cook(Tomato)
Cook(Salad)
Frozen Lake
Horizon
60
60
30
50
20
Table 3: Maximize horizon of environments.
we use 11 tasks, such "examine the laptop with the desk lamp." There are no auxiliary rewards, except for a reward 1.0 for reaching the final goal.
Examples of observations and admissible action descriptions for three environments are shown below:
For )LFk)rld(Pick)
Current observation: You are in the middle of a room. Looking quickly around you, you a armchair I, bed diningtable 1, drawer 2, drawer 1, garbagecan 1, sidetable 2, and a sidetable 1.
Observation:
Your task to: put some cellphone on armchair..
Admissible actions You are allowed to take the following actions: go to armchair 1, go to bed 1, go to diningtable go drawer go drawer go garbagecan 1, go to sidetable 1. go to sidetable 2, inventory, look.
For Overcooked( (Tomato)
Observation: Your task is to serve the dish a bowl only containing chopped tomato. There is fixed cutting board in the room. You notice tomato the table. Currently you don't 't have anything in hand. Admissible actions You allowed to take the following actions: pick up the tomato, take the bowl, walk to the cutting board, serve nothing, chop nothing.
For Frozen Lake
Observation: You are in a 4x4 grid where each cell is either frozen cell (F) indicating safe position, You must avoid reach a hole (H). and the goal (G) that we aim to reach.
Currently the agent is at position (0, 0), which direction need to take? Admissible actions: Go Left leads to position (0, 0), start position. Go Right leads to position (0, 1), frozen cell (F). Go Up leads to position (0, 0), start position. Go Down leads to position (1, 0), a frozen cell (F).
Input:
Admissible actions You are allowed to take the following actions: Go Left, Go Right, Go
Up. Go Down.
9.2 LLM PRIOR IMPLEMENTATION
There are two ways to sample an action from the LLM prior PLLM(-)s:). First, since the state and action can be described as text, and assuming the action a consists of le tokens, the probability of the
LLM generating action a is given by p(a)s;) II PLLM(a, St, azz). Based on this probability, the first type of LLM prior computes distribution actions, denoted dist (al (s)) E., exp exp p(al)s.)
exp p(a) ex)
In contrast, the second approach involves sampling free-form output from the I.I.,, which is then mapped to executable action via simple rule-based projection P. We denote this LLM prior, which relies on mapping the LLM's output, as PLIM where a ~ PLIM (-1st) ++ o LLM(-Is:)) = P(o).
We observed a performance gap between these two types of LLM action priors. For instance, in the ALFWorld pick & place task, PL.1M achieved a winning rate of 19%. while dist failed, with a map winning rate of 0%, both using the LLM Qwenl.5-7B. Therefore, in this work, implement PLLM as the LLM prior for thoughtfully exploiting the power of the LLM in SDM tasks, which we denote PLLM for simplicity. An Example of L.M''s output is given as:
19

---

Published as a conference paper at ICLR 2025
For ALFWorld(Pick) from Qwen- 1.5 7B
Input: Current observation: You are in the middle of room. Looking quickly around you, you see a armchair 1, bed 1, diningtable 1, a drawer 2, a drawer 1, a garbagecan 1, sidetable 2, and a sidetable 1.
Your task is to: put some cellphone
armchair.. You are allowed to take the following actions: go to armchair I, go to bed I, go to diningtable 1, go to drawer 1, go to drawer 2, go to garbagecan 1, go to sidetable 1, go to sidetable 2, inventory, look. Please select an action from the admissible actions.
LLM Output to Action: 1: take ed from sidetable 2, -->
take ed 2 from sidetable 2
2: Open drawer 2 Result: You open door, -- Open drawer
3: look Description: You look around the room.,
look
look > Current observation: You arrive at loc, - look
5: open drawer Loc 8: You open drawer 2 Loc 8-> open drawer Loc
We map the LLM's output to an executable action using simple rule-based method, extracting the executable actions directly from the LLM's output.
9.3 PAPEE-PARAMETERS
Our algorithms trained one machine with 40G A100. based Pytorch-GPU 2.1 and cuda12.4. Table algorithms, set the hyperparameter for regulating the Q-values, shown in Eq. 5.0.
reports the main hyper-parameted of our algorithms. For all CQL-based
Offline Datasets The composition of offline datasets is shown in Table which reports the number of (s, s') examples.
Table 4: The hyperparameters ALFWorld(Pick)
Baselines DQN-Prior
Learning Rate Epochs
Batch Size
Update Frequency
LLM
5e-4
128
0.01
CQL-Prior
5e-4
4
128
5
0.01
GFlan-Prin
le-4
16
64
16
Flan-T5 Small
0.01
Table 5: The byperparameters on ALFWorld(Examine)
Baselines
Learning Rate
Epochs
Batch Size
Update Frequency
L.M
Cx
DON-Prior
5e-4
4
128
10
0.01
CQL-Prior
5e-4
128
10
0.01
GFlan-Prior
le-4
16
64
16
Flan-T5 Small
0.01
Table 6: The hyperparameters on Overcooked(Tomato
Baselines
Learning Rate
Batch Size
Update Frequency
LLM
DQN-Prior
5e-4
128
0.1
CQL-Prior
5e-4
128
0.1
GFlan-Prit
le-4
64
16
Flan-Tr T5 Small
0.1
Table 7: The hyperparameters on Overcooked(Salad)
Baselines
Learning Rate
Batch Size
Update Frequency
LLM
a
DON-Prior
5e-4
128
0.1
CQL-Prior
5e-4
128
0.1
-
le-4
64
16
Flan-T5 Small
0.1
20

---

Published as a conference paper at ICLR 2025
Table 8: The hyperparameters Frozen Lake
Baselines
Learning Rate
Batch Size
Update Frequency LLM
DON-Prior
5e-4
128
0.1
Table 9: The hyperparameters on ALFWorld((ick)
Baselines
Learning Rate
Batch Size
Update Frequency
LLM
DQN-Prior
5c-4
128
10
0.01
CQL-Prior
5e-4
128
10
0.01
GFlan-Prion
le-4
64
16
Flan-T5 Small
0.01
Table 10: The hyperparameters on ALFWorld(Examine)
Baselines
Learning Rate
Batch Size
Update Frequency
LLM
DQN-Prio
5e-4
128
10
0.01
CQL-Prior
5e-4
128
10
0.01
GFlan-Prior
le-4
64
16
Flan-T5 Small
0.01
Dataset ALFWorld(Pick)
Total Examples
Good Examples
Bad Examples
6572
6572
0
Salad( 1000)
1186
574
612
Salad(4000)
3892
2209
1683
(alad(8000)
7571
4224
3347
Salad(1 2000)
11939
4004
7935
Salad(24000)
23869
7191
16678
Table 11: The composition of the offline dataset
9.4 ADDITIONAL RESULTS
This section presents additional experiments to further demonstrate the effectiveness of our method, including: 1) Our approach, only relying on I.LM querying, more time-efficient compared to the LLM fine-tuning-basee approach. 2) Our approach is adaptable partially observed BabyAI environments. 3) The drawback frequently querying LLMs for action priors can mitigated through eaching techniques. 4) Inspired by the success of value-based DQN-Prior, conducting policy-based RL exclusively within the prior space is also feasible and significantly improves sample efficiency.
(PEc)
Overcodked (Tbmato)
Overrocked (Salad)
0.8
DON--Prior
1.0
1.0
0.8
Gilan
LLM-Prior
0.6
Reward
0.4
0FlinsProrttee 1)
0.4
0.2
0.0
-0.5
800
1000
150
Egisodes
Fpisodes
Episodas
Figure 5: The ablation of the number of action proposals & used for approximating KL divergence between the required action policy and the LLM prior action distribution.
Additional ablation on i for GFLan-Prior As shown Figure 5 when we sample k: = 5 action proposals from the LLM to approximate the KL divergence, performs better than when using A 1. This is because, with more proposals, the optimal action is more likely to included in the subset, thus better guiding the action policy.
Additional Ablation Study on the Caching Technique for Time Efficiency Most studies on RL fine-tuning for LLMS Tan al. (2024); Christianos al. (2023): Wen al. (2024) face the
21

---

Published as a conference paper at ICLR 2025
Overcooked (Tomato)
1.00
0.75
Reward
0.50
DON-Prior(per step) (acch)))r(Cacch)
0.25
0.00
TWOSOME
-0.25
0
100 200 300 400 500 600 Episodes
Figure 6: Comparison of DON-Prior(Cachh), DQN- LLM per step) and TWOSOME on Overcooked (Tomato).
challenge of querying the LLM for each step. They use large LLMs (up to 7B parameters) as the action policy, and require querying the L.M at each step to collect data for online training. Since the I.LM prior is fixed, we can resort to the caching technique within our framework to alleviate the need for querying the L.M at every step. We introduce a variant of the DQN-Prior algorithm, called DQN-Prior (Cache), which uses cached action candidates for previously seen observations and queries the LLM prior, storing its outputs only for new observations. We present the time and computational costs of our DON-Prior(quer LLM each time) and DQN-Prior (Cache), well the RLFT baseline, TWOSOME Tan et al. (2024), for learning to converge on Overcooked (Tomato) in Table 12 Results of baselines are shown in Figure 6 TWOSOME fine-tunes the LLM (Qwen-1.5 7B for a fair comparison) using the LoRA technique Hu et al. (2021) and the PPO algorithm. The time- consuming of our MQN-Prior(Cache) is significantly lower than the TWOSOME, while maintaining comparable performance. The caching technique is well-suited to our framework, as the 7B LLM capable providing reliable sub-action space. DQN-Prio (Cache) significantly reduces query time costs but also achieves superior performance compared to DQN-Prior (query LLM per step), as frequent queries introduce more uncertainty.
Table 12:
The training costs of DQN- Prior (Cache). DQN-Prior (query LLM per step), and TWO- SOME on Overcooked (Tomato) are reported for one training seed.
Baselines
Training Time(min) GPU Usage
Query LLM(Qwen-1.5 7B) Times
DON-Prior(per step)
31.1
28298 MB
2933
DON-Prior(Cachor
13.8
28340 MB
434
TWOSOME
60.4
25014 MB
5000
Additional Experiment on BabyAl We compare our policy- algorithm, GFlan-Prior, with GFlan the text-based BabyAl environment Carta al. (2023), which features partial observability and random map initialization for each episode. We focus the BabyAl task 'Go to the red ball' and set the maximum horizon at 30. The comparison results are presented in Figure 7 Our GFlan-Prion which integrates both LLM prior knowledge and environmental interaction experience, successfully accommodates the partially observed BabyAI and outperforms both GFlan and the LLMM-rror.
Ablation Study on Using a Large LLM for the Q-Network We conducted an experiment on (vereooked(Tomato) compare the effects of using an LLM as the action prior versus using it to encode state actioenttonnnnn pairs. We evaluated three baselines:
DQN-Prior: Uses Qwen- 7B to generate sub-action space and BERT to encode (s, a)
pairs for Q-value estimation.
DQN (BERT): Uses BERT to encode (s, a) pairs and performs DQN across the full action space.
DQN (Qwen-1.5 7B): Uses 1wen-1.5 7B to encode (s, a) pairs and performs DQN across the full action space.
22

---

Published as a conference paper at ICLR 2025
Babyal
Babyal
1.0
Gflan
Normalized Reward
Normaltzed Reward
0.8
1IM
0.6
0.4
0.4
0.2
Gilan LLM.
100
200
300
400
500
250 500
750 1000 1250 1500 1750 2000 Episodes
Episodes
(a) Max Horizon as
30
(b) Max Horizon as 100
Figure The comparison of policy-base algorithms on BabyAl.
Overcooked (Tomato)
1.00
0.75
Reward
0.50
0.25
DQN(Bert)
DQN-Prior
0.00
NQN(Qwen-1.5 7B)
-0.25
50
100
150
Episodes
Figure 8: An ablation to verify the effectiveness of using the LLM to provide prior actions versus to encode state-action pairs.
To reduce the overhead of querying the LLM for every state, we store its outputs: For DQN-Prior, we store action proposals. For DQN (Qwen-1.5 7B), store (s, a) embeddings.
As shown in the Figure and Table 13 DQN-Prior outperforms both DQN (Qwen-le 1.5 7B) and DQN (BERT). By leveraging the LLM's high-level reasoning, DQN-Prior narrows the exploration space to a reliable action sub-space, improving sample efficiency. While DQN (Qwen-l. 1.5 7B) benefits from the 7B LLM's superior semantic understanding and outperforms DQN (BERT), (s,a) embeddings cannot be directly used for action selection. Instead, mapping from embeddings to Q-values must be learned from scratch to correctly rank the entire action space, thereby increasing sample complexity.
In addition, DQN-Prior requires less time and CPU memory compared to DQN (Qwen-1.5 7B). The higher memory and time consumption of DQN (Qwen-1.5 7B) are primarily due its large embedding size (4096 dimensions), which significantly increases the cost of training a mapping (MLP adaptor) from action embeddings Q-values, as well as the memory required to store these embeddings.
Table 13: The training costs of DQN-based baselines: DQN-Prior, DQN(BERT), DQN(Qwen1..54)) 1.5-7B.
Both DQN (BERT) and DON (Qwen-1.5 7B) involve the of the 7B L.M.. To minimize the overhead of querying the LLM for every state, we pre-store its outputs: for DON-Prior, we store action proposals, while for DQN (Qwen-155 7B), store (s, a) embeddings.
Baselines
Training Time(min)
GPU Usage CPU usage
Embedding Dimension
DQN- Prior
13.8
28340 MB
18KB
768
DON(BERT)
16.2
17864MB
768
DON(Qwen-1.) 7B)
74.8
22662MB
22724MB
4096
23

---

Published as a conference paper at ICLR 2025
Applying Direct Posterior Sampling on the Top of PPO Inspired by the success of direct posterior sampling in value- based reinforcement learning (RL), which effectively explores and updates the value function within the LLM prior action space, recognize that LLMs provide narrow and reliable sub-action space. Building on this success, we aim to investigate policy-based RI. within this action space. Specifically, we seek to learn a policy that approximates the soft Q-distribution over the LLM prior space:
Q(s, a)
(a(ass)
Exrect Q(s, d')
In practice we directly apply Proximal Policy Optimization (PPO) within the prior sub-action space. Formally, for the i-th iteration, the policy To, learned by optimizing:
To (T( (st)
)
arg max (s))(s))pzzan)) (o, CLIP
Å. + (7(((T(-18=)) (26)
8
To,-, (ae(s))
where A, is the estimated advantage, y) is the hyperparameter that encourages exploration. In the textual environments we considered, we used GFlan Carta et al. (2023) to train PPO with Flan-T5 small as the action policy, For simplicity, we denote the policy-based algorithm as (FlaanPriior(Sub Action Space) or GFlan-Prior(Sub AS).
The comparison of pasicy-based RL with LLM prior shown in Figure (rraan)rioo(KL)) indicates the PPO algorithm with the incorporation of LLM prior through the addition of a KL-penalty. These results demonstrate that incorporating LLM prior into PPO training, either by adding a KL-penalty by providing sub-optimal sub-action space, can improve the sample efficiency of RL training. Additionally, while both GLL and PrrlannPrioo(Sute AS) are able to successfully complete the final task with rewards approaching 1, GFlan- GFlan-Prior(Suul AS) achieves higher rewards due to its explicit exploration and optimization within the sub-action space, which is more efficient than the implicit use of LLM prior by adding KL penalty.
Overcooked (Tomato)
1.2
1.0
Reward
0.8
0.6
GFlan
LLM-Prior
0.4
GFlan-Prior(KL)
0.2
GFlan-Prior(Sub AS)
0.0
50
100
150
Episodes
Figure 9: The comparison of policy-based algorithms with LLM prior.
10 ADDITIONAL RELATED WORKS
Value-based Inference of LLM Not limited to the pre-training phase, it has been recently inves- tigated how scaling law applies at inference time for LLM. This often referred to as searching over an enormously large output space with guidance from value estimation. For example, the naive case beam search (citation) which selects the top K possible sequence based on cumulative probabilities as opposed to greedy search. The value estimate here is the likelihood of the sentence as predicted by the language model. (Wang et al. 2022) propose CoT-SC, which searches over reasoning paths and selects the most frequent answer. Tree-of-Though (ToT, (Yao et al.) 2023) adopts depth/breadth-rfe
search. (Hao et al.) introduce Reasoning- -iiooonnnn-------eee via-Planning (RAP) which uses Monte-Carlo Tree Search and value estimate from prompting LLM. TS TS-LLM designed by (?) guides MCTS with a learned value function conditioned on state and a learned Outcome- -upperiiee Reward Model (ORM). (Li et al. 2024) proposes Q-probing adapt pre-trained language model to also
24

---

Published as a conference paper at ICLR 2025
maximize tasks-specific reward function in code generation tasks. Furthermore, recent works investigate Process-supervised Reward Model (PRM, (Lightman et al. (003)) and apply in LLM inference, (McAleese et al. 2024) searches over step-wise critic model in code reviewing tasks. (Wang et al. 2024a) learn a PRM in math reasoning and use it for LLM fine-tuning, (Snell et al.
2024) provides a detailed benchmark of different inference methods. (Zhang al. (024a) represents the score as the probability of a single text token (e.g. 'Yes' or 'No") under the context and the prompt. Even though our proposed method can be seen an example of value-based inference, we focus on a different perspective where we investigate how LLM helps conventional value-bar basee reinforcement
learning algorithms in complex environments.
LLM-based Reward Function for RL Beyond serving high-level planners for SDM tasks, LLMs are also utilized for generating reward signals, which are subsequently used for RL training. For example, (Kwon et al. 2023) leverages LLMs directly generate reward signals prompted with historical interaction information. Motif (Klissarov et al., 2023) employs LLMs to annotate pairwise preference datasets, which are then used to train reward function. Given the challenges LLMs face in directly generating low-level actions for robotics tasks with continuous representations, several studies Yu et al. 2023 Ma et al.) 2024) have instead utilized the coding capabilities of GPT-4 to generate human-lever reward code. These reward codes are then employed to facilitate RL training.
More discussion can be found in Section 2
25


Reviews:
[{"review_id": "ktRP1YNfqG", "reviewer": "Reviewer_h6X1", "summary": "This paper considers treating LLMs as prior action distributions and integrating them into RL frameworks. Specifically, authors proposes to use LLM for exploration,  Q target calculation and imposing conservatism. Experiments on the ALFWorld and Overcooked environments illustrates the effectiveness of the proposed framework.", "questions": "1. The experiments are only conducted on simple environments with discrete action spaces. What if continuous action space (e.g., offline RL on the D4RL benchmark)?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper is clearly written and easy to follow.\n2. The considered \"incorporating the background knowledge of LLM into RL\" is an interesting topic.", "weaknesses": "1. There are no convergence guarantees of sampling using LLM as prior for Exploration and Q-function update. Moreover, there is no guarantee that learning the Q-function as shown in Equation (7) will give us a conservative Q-function.\n2. If the LLM can not provide a good prior, constraining the learned policy as shown in Equation (8) will result in sub-optimal policies.\n3. The RL process has to query the LLM each time we update the policy, which is time-consuming and computation-consuming.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "mSaNyt7fOw", "reviewer": "Reviewer_vRpg", "summary": "In this work, the authors propose idea of using large language models (LLMs) for providing prior information about execution of a task, which simplifies exploration in reinforcement learning (RL) and improves sample efficiency. The core idea is to treat LLMs as prior distribution over actions and then apply variational inference in conjunction with posterior sampling. The authors demonstrate this technique on both value-based and policy-based approaches. In value-based implementation, a subset of the action space is chosen for Q-learning instead of using entire action space. This subset selection is done using LLM prior. Similarly, in the policy-based implementation, a KL-penalty is introduced with respect to the prior LLM action distribution, but implemented differently from regular KL penalty in RL from human feedback (RLHF), by posterior sampling of multiple candidate actions. The experiments are conducted on text-based RL environments with large action spaces and multi-step horizon (note this is crucial for demonstrating the difference against RLHF techniques). Value-based baselines include DQN for online learning and CQL for offline learning. For policy-based baseline, PPO algorithm is compared against.\n\nThe results demonstrate the superiority of the approach in (1) overall returns, (2) generalizing beyond seen tasks, (3) sample efficiency.", "questions": "Questions and Suggestions:\n\n- Should the proportionality sign on line 147 be equal to? I understand the sentence is correct in its current form.\n\n- The $p$ is overloaded on lines 112-113, especially, please change the \"$p_{LLM}$ as the action prior $p$\" sentence where the latter $p$ is said to be not confused with $p$ used for probability of start state or action given a state in the definition of probability of trajectory.\n\n- I am not sure by what is meant by \"Denote that the above …\" in Proposition 1. Please rewrite this for clarification.\n\n- Line 215 - \"reply buffer\" should be changed to \"replay buffer\" \n\n- Is the gamma in equation 8 the same as the discount factor?\n\n- Line 323 - \"avoiding full into holes\". Please rectify this.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is written very clearly with appropriate figures, well-formatted equations and more importantly lucid logical flow. The contributions of the paper are specified along with particular sections corresponding to individual contribution. Preliminaries provide enough details and are not overexplained. \n\nThe core idea of the paper is pretty neat -- LLMs are trained on large amount of data and they have knowledge about tasks if not fine-grained controllability; this knowledge is useful to guide an RL training. \n\nI like how the value-based implementation is designed -- sampling a set of actions using LLM prior, taking Q estimates, sampling actions using their Q estimates. The method is well grounded in the theory where direct posterior sampling is shown to be proportional to LLM prior multiplied with soft-Q distribution. The idea being simple, can be plugged with both online and offline Q-learning approaches.\n\nThe versatility of the idea is shown by applying it to policy-based RL. The LLM prior guided learning takes form similar to KL-regularized RL, hence modifying PPO in a straightforward way would lead to easy access to prior knowledge. However, I do I have some contentions around the novelty of this part which is discussed in weakness section.\n\n\nThe choice of the experimental setup is also very apt. The environments are chosen so that LLMs can provide textual action descriptions which then can be executed through text to action conversion in the environment's processing.", "weaknesses": "Practicality of value based approaches: Although the value-based direct posterior sampling technique described in the paper is logically sound, I wonder about its practicality. Mainly, value-based approaches when applied to combinatorically large action spaces like text generation is practically infeasible. Hence, the paper's proposed method might be hard to apply beyond toy textual environments. Noticeably, these environments are designed to accept actions in a particular format or use rigorous action decoding, etc. to deal with the combinatoric complexity. Such environments are rare in reality.\n\nProblems with choice of prior inducing LLMs: Two LLMs might produce semantically same action but different text. I am curious to know whether the Q-learning based approaches would assign similar values to both. In general, I feel that prior inducing LLM and its compatibility with the environment's action decoding plays major role when they are used to train RL policies in text-based sequential decision making. \n\nNovelty of the additional KL penalty in policy-based RL: ELBO-based RL approaches [1, 2] have been already established in the RL literature. The additional KL-penalty proposed in the current paper is based on the same. In single step decision making, the KL-penalty on an LLM prior would be exactly same as KL-penalty used along with supervised fine-tuned (SFT) policy. Thus, in that setting, the paper's technique lacks novelty. Now, coming to the multi-step decision making, unless the proposed direct posterior sampling is also applied on top of PPO, additional KL-penalty is equivalent to using a good enough prior policy (e.g. behavior cloned policy). Therefore, the novelty of the paper is limited in policy-based RL.\n\nCombined together,  I find that paper provides great insights into a direct posterior sampling paradigm for using prior knowledge present in the LLMs, but the approach might have limited applicability beyond hand-designed textual environments, its applicability is limited by limitations of value-based approaches, and its novelty in policy-based RL is not significant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "2n3jjWYOFy", "reviewer": "Reviewer_8Lkv", "summary": "Authors present a new way to train LLM-based agents in interactive sequential environments. Their approach is based on leveraging a potentially large frozen LLM to generate action subsets for encountered states. The action subset is then used to simplify the RL finetuning of a smaller LLM. Their approach can be adapted to value-based (giving a subset of actions for the LLM Q-network) and policy-based scenarios (constraining the LLM policy with a KL-div loss over the subset of actions). Authors showcase how their approach outperforms baselines in TextWorld and Overcooked, both in online and offline learning settings.", "questions": "l.286 \"We plot the rewards averaged over the final third of the training processes for trainable baselines\" --> Evaluation of a model performance is not performed by conducting periodic performance measurements over a test set of tasks ? Could you elaborate on how performance is measured ?\n\nIn your policy-based version, which you present in section 3.3 as an extension of the work of Carta et al., 2023, I am not sure to understand what is trained is what is fixed. Are you, like in Carta et al, training all parameters of the LLM ? What is being trained by your modified PPO loss ? Would be good to have the policy-based version of figure 1 somewhere. My understanding after reading section 4.2 is that you have a large fixed pretrained model used to compute action priors, and you train another smaller Flan-T5 as the policy.\n\nIn your offline learning experiments, by looking at table 1 it looks like you are applying a vanilla DQN to offline datasets. It comes with the risk of overestimating values of actions that are not supported by the dataset. Would it make more sense to use value-based methods designed for offline learning, e.g. Implicit Q Learning from Kostrikov et al ?\n\nIn your offline learning results, I see that CQL underperforms wrt Behavioral Cloning. This is unusual: any ideas on why ?\n\n\n\n\nGiven the aforementioned questions and potential limitations of this paper, I will go for a weak reject, and look forward to the discussion to update my score.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "This paper tackles a very important topic: how to leverage LLMs into sequential decision-making tasks incorporating language.\nI enjoyed reading it.\nPerformance improvements over considered baselines are impressive.", "weaknesses": "### Value-based version\n\nI am having trouble understanding how useful the value-based formulation is. What I mean is that, from my understanding, at inference time you need both the frozen LLM and the trained Q-network (which is a smaller LLM). This makes the approach quite expensive regarding memory. Regarding inference, it is also costly: the frozen LLM must provide k outputs before being able to use the Q-network LLM, for each step in the environment.\nA good comparative study here to motivate this approach could be to have a baseline which discards the frozen-LLM. What happens in this scenario ? Is it what you call the DQN baseline ? Here, it would be interesting to see what happens if you use a larger LLM during training (something which induces a memory burden comparable to using a frozen LLM + a smaller trained Q-net LLM).\n\n### Comparative study with GFlan from carta et al\n\nThe only baseline considered from the litterature is GFlan from carta et al. While I am not expert enough to assess whether this is sufficient, I have concerns about this comparative study. \nIn carta et al, GFlan is evaluated on a text-based version of the BabyAI gridworld environment:\n\nWhy not featuring this environment in the study ?\n\nAny intuitions on why GFlan is failing completely in ALFWorld (TextWorld) while it does maintain some performances in Overcooked ?\n\n\n\n\n\n\n### Clarity\n\nWhile the paper is mostly clear, I had trouble understanding the exact training and inference pipeline. The paper could benefit an update of figure 1 (along with providing the policy-based version of figure 1 somewhere).The right side of Figure 1 could be improved: in the current version, it is hard to visually understand that the Q function is used to weight each input actions and that these weights are used for proportional sampling of the next action. Also in figure 1 it is hard to understand that there are two pre-trained LLMs being used, one frozen generating an action subset, and a smaller one trained and used as a critic. I only understood that there was two pretrained LLMs being used in the DQN-Prior baseline at line 395 \"... by combining the Q-function, trained with Qwen-1.5 7B\". I would suggest making it clear early on.\n\nl.97-99 \"we first generate a free-form output from the LLM, for example, a 7B LLM, which is then mapped to an executable action through a simple rule-based projection.\" \n--> I would recommend moving at least part of what is discussed in the appendix into the main body of the paper, to be clearer about this. \nRelated to this: I only understood that you repeat this output-to-action mapping multiple times at line 166-167. It would be clearer to mention it sooner in the paper.\n\n\n### Minor / Typos\n\nThere are 3 acronym definitions in the abstract, which is a lot. Consider reducing this number, e.g. dropping SDM which is used only once.\n\nl31. \"human-AI dialogue (McTear, 2022; Li et al., 2019),\" --> make sure to sort citations in ascending order wrt year. This error appears multiple times.\n\nl. 232 --> (Snell et al.) missing year in citation.\n \nl.104 \"by (Levine, 2018)\" --> use \\cite not \\citep when directly mentioning the citation in the sentence.\n\nl.304 \"can be found in the Appendix\" -->  refer to the appropriate appendix section(s) rather than a general reference\n\nl.323 \"while avoiding full\" --> falling\n\nl.352 \"Behavior Clone(BC).\" --> Behavioral Cloning ?\n\nResults-related figures could be moved closer to the result analysis in section 4.3 to simplify reading\n\nl.404 \"k = 1 indicates the performance of the LLM on its own, without involving the Q-function\" \n--> Why calling this scenario k=1 and not k=0 ? you could theoretically use your frozen LLM to create a single action subset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "Hd57aeJnhs", "reviewer": "Reviewer_uGBu", "summary": "This paper proposes an approach to leverage LLM priors for RL, both in the online and offline settings.\nThe setting is restricted to text-based observations. Specifically, they estimate a prior p(a|s) by sampling K actions from the LLM conditioned on the state, and use these to either restrict the Q-values to the actions with non-zero probability (for online or offline Q-learning based methods), or to enforce a KL penalty with respect to the policy’s distribution over actions. They evaluate on several text-based environments and show improvements in terms of sample efficiency over vanilla RL. Overall, the paper does a decent job of illustrating the benefits of the method within its scope, but the scope itself is quite limited (in terms of broad applicability and scalability - more on this below). Also, there are a number of alternate methods for incorporating LLM priors into RL systems, that are not discussed or compared to.\n\n=====\nUpdate after rebuttal: the authors have added several additional experiments which partially address the initial limitations I discussed. Although I still think the environments are on the simple side, and the method in its current form is difficult to scale to high-throughput settings, the results are convincing enough in demonstrating the approach's merits and give hints of wider applicability to meet the bar for acceptance.", "questions": "If the authors can address any number of the points raised in the Weaknesses section, I would be willing to raise my score accordingly. I also have a few minor comments below:\n\n- In Section 4.1, it would be helpful to give some more detail on the environments you use. The current descriptions are quite high level and not everyone may be familiar with these environments or have intuitions about them. Maybe include a figure with example observations/actions for each one - you still have half a page under the page limit so there should be space. \n- Some of the references are missing dates, for example (Snell et al), line 232. \n- This is a bit of a nitpick, but in Figure 1, is there a reason the OpenAI logo is on the LLM box? My understanding is that this paper does not use OpenAI’s LLMs, and instead uses open models (Qwen and Llama). I don’t think an institutional logo is necessary, but if you do want to include one, it would seem more fair to give credit to those that made the models that you use available.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the method is simple and makes intuitive sense. Its benefits are clearly demonstrated, although in fairly simple text-based environments. Although other works have proposed ways to incorporate LLM priors into RL systems, to my knowledge this particular approach is new. \n- the paper is overall clearly written and easy to understand. I have a few suggestions below, but they are relatively minor.", "weaknesses": "- the method is limited to settings where the observation consists of text, which is a big limitation. One could imagine workarounds to apply this to images, where one first uses a pretrained captioning model to extract textual captions and then feed these to the LLM. However, it’s not clear if this would give enough granularity to produce a reasonable action prior and this would need to be checked experimentally. \n- the method, as I understand it, will have difficulty scaling to high-throughput RL settings because it requires a call to the LLM for each observation encountered by the agent (which is computationally expensive). This also limits its applicability to small tasks or offline settings where the LLM outputs can be precomputed in parallel and cached. Are there any ways to address this? For example, could you distill the LLM action prior into a lightweight model and use this instead?\n- there are other ways to incorporate LLM priors into RL systems, namely through the reward function rather than the action distribution. See for example the references below [1, 2, 3, 4] (this is not a complete list). It is currently unclear what the pros/cons of the proposed method are compared to this other family of approaches. It would be helpful to include at least one method from this family as a comparison in the experiments. Also, these are not discussed in the related work section.  \n\n\n[1] Eureka: https://arxiv.org/abs/2310.12931\n[2] Language2Reward: https://arxiv.org/abs/2306.08647\n[3] Motif: https://arxiv.org/abs/2310.00166\n[4] Reward Design with LLMs: https://arxiv.org/abs/2303.00001", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)