Zero-shot Model-based Reinforcement Learning using Large Language Models
Authors: Abdelhakim Benechehab
Youssef Attia El Hili
Ambroise Odonnat
Oussama Zekri
Albert Thomas
Giuseppe Paolo
Maurizio Filippone
Ievgen Redko
Balázs Kégl
Venue: ICLR 2025 Poster (2025)
Keywords: Model-based Reinforcement Learning
Large language models
Zero-shot Learning
In-context Learning
Primary area: learning on time series and dynamical systems
TL;DR: We investigate in-context learning using LLMs for dynamics learning of Markov Decision Processes with applications in off-policy reinforcement learning.

Abstract: The emerging zero-shot capabilities of Large Language Models (LLMs) have led to their applications in areas extending well beyond natural language processing tasks. 
In reinforcement learning, while LLMs have been extensively used in text-based environments, their integration with continuous state spaces remains understudied. 
In this paper, we investigate how pre-trained LLMs can be leveraged to predict in context the dynamics of continuous Markov decision processes. 
We identify handling multivariate data and incorporating the control signal as key challenges that limit the potential of LLMs' deployment in this setup and propose Disentangled In-Context Learning (DICL) to address them.
We present proof-of-concept applications in two reinforcement learning settings: model-based policy evaluation and data-augmented off-policy reinforcement learning, supported by theoretical analysis of the proposed methods.
Our experiments further demonstrate that our approach produces well-calibrated uncertainty estimates. We release the code at https://github.com/abenechehab/dicl.

Published as a conference paper at ICLR 2025
ZERO-SHOT MODEL-BASED REINFORCEMENT LEARN- ING USING LARGE LANGUAGE MODELS
Abdelhakim Benechehabl 12, Youssef Attia EI Hili', Ambroise Odonnat! Oussama Zekri
Albert Thomas', Giuseppe Paolo', Maurizio Filippone', levgen Redko', Balázs Kégl'
Huawei Noah's Ark Lab, Paris, France
Department of Data Science, EURECOM
3 Inria, Univ. Rennes 2, CNRS, IRISA
4 ENS Paris-Saclay
5 Statistics Program, KAUST
ABSTRACT
The emerging zero-shot capabilities Large Language Models (LLMs) have led to their applications in areas extending well beyond natural language process- ing tasks. In reinforcement learning, while L.LMS have been extensively used text-based environments, their integration with continuous state spaces remains understudied. this paper, we investigate how pre-trained LLMs can be lever- aged to predict in context the dynamics of continuous Markov decision processes. We identify handling multivariate data and incorporating the control signal as key challenges that limit the potential of LLMs' deployment in this setup and propose Disentangled In-Context Learning (DICL) to address them. We present proof-of- concept applications in reinforcement learning settings: model-based policy evaluation and tuta-augmented off- off-policy reinforcement learning, supported by theoretical analysis of the proposed methods. Our experiments further demon- strate that our approach produces well-calibrated uncertainty estimates. We re- lease the code at https::/github..om/aaaa http://www....bbbbeecchhhhhhhe..
for each component:
concat
Trajectory
LLM
Prediction
[x(c
R (T-e)
(p: e.g. PCA
(br F Frea
R 12x(( +
c: components
h: prediction horizon
auloregressive
Figure 1: The DICL Framework. DICL projects trajectories into a disentangled feature space before performing zero-shot forecasting using pre-trained LLM and in-context learning.
INTRODUCTION
The rise of large language models (LLMs) has significantly impacted the field of Natural Language Processing (NLP). LLMs (Brown et al., 2020; Hugo Touvron the Llama 3 team., 2024), which based on the transformer architecture (Vaswani et al., 2017), have redefined tasks such as machine translation (Brown et al., 2020), sentiment analysis (Zhang et al., 2023bb, and question answering (Roberts et al., 2020; Pourkamali & Sharifi, 2024) by enabling machines to understand and generate human-like text with remarkable fluency. One of the most in- triguing aspects of LLMss is their emerging capabilities, particularly in-context learning (ICL) (von Oswald et al., 2023). Through ICL, an LLM can learn to perform a new task simply by being pro- vided examples of the task within its input context, without any gradient-based optimization. This
Llama 2 team., 2023; Dubey &
1 "ooresponndence ahdelhakim. www.hihhhhhccchhhhb.aboooo gmail.com. Work done while Huawei Noah's Ark Lab.

---

Published as a conference paper at ICLR 2025
phenomenon has been observed not only in text generation but also in tasks such image classifi- cation (Abdelhamed et al., 2024; Zheng et al., 2024) and even solving logic puzzles (Giadikiaroglou et al., 2024), which is unexpected in the context of the standard statistical learning theory. To our knowledge, ICL capabilities of pre-trained L.LMs have been only scarcely explored in reinforce- ment learning (Wang et al., 2023) despite the demonstrated success of the former in understanding the behavior of deterministic and chaotic dynamical systems (Liu et al., 2024c).
In this paper, we show how ICL with pre-trained LLMs can improve the sample efficiency of Re- inforcement Learning (RL), with two proof-of-concepts in policy evaluation and data-augmented off-policy RL. Following the dynamical system perspective on ICL introduced in Li al. (2023) and experimentally studied in Liu et al. (2024c), we use the observed trajectories of a given agent to predict future state and reward in commonly used RI. environments. To achieve this, we solve two crucial challenges related to considering continuous state-space Markov Decision Processes (MDP): 1) incorporating the action information into the L.M''s context and 2) handling the interdependence between the state-actions dimensions, as prior approaches were known treat multivariate data's lata's covariates independently. Our framework, DICL (Disentangled In-Context Learning), is summa- rized in Fig. 1. The core idea of DICL is to apply feature space transformation, denoted as ip, which captures the interdependencies between state and action features in order to disentangle each dimension. Subsequently, Large Language Model (LLM) is employed to forecast each compo- nent independently in a zero-shot manner through in-context learning. Finally, the predictions are transformed back to the original trajectory space using the inverse transformation yp
Our approach leads to several novel insights and contributions, which we summarize as follows:
1. Methodological. We develop novel approach to integrate state dimension interdepen- dence and action information into in-context trajectories. This approach, termed Disen- tangled In-Context Learning (DICL), leads new methodology for applying ICL in RL environments with continuous state spaces. We validate our proposed approach on tasks involving proprioceptive control.
2. Theoretical. We theoretically analyze the policy evaluation algorithm resulting from multi- branch rollouts with the LLM-based dynamics model, leading to a novel return bound.
3. Experimental. We show how the LLM's MDP modeling ability benefit two RL appli- cations: policy evaluation and data-augmented offline RL. Furthermore, we show that the LLM is calibrated uncertainty estimator, a desirable property for MBRL algorithms.
Organization of the paper. The paper is structured as follows: Section 2 introduces the main concepts from the literature used in our work (while a more detailed related work is deferred to Appendix B). We then start our analysis in Section 3.1, by analyzing LLM's attention matrices. DICL is presented in Section 3.3, while Section contains different applications of the proposed method in RL, along with the corresponding theoretical analysis. Finally, Section 5 provides short discussion and future research directions triggered by our approach.
BACKGROUND KNOWLEDGE
Reinforcement Learning (RL). The standard framework of RL is the anfinite-horizon Markov de- cision process (MDP) M = (S. A, P.r, No. 7) where represents the state space, AA the action space, P S A + S the (possibly stochastic) transition dynamics, r S x AA + R the reward function, 14 the initial state distribution, and 7 E [0, 1] the discount factor. The goal of RL is to find, for each state E S, distribution क (s) over the action space A, called the policy, that maxi- mizes the expected sum of discounted rewards nTTT) == E P: (2 7'r(s,, az)]. Under a policy n, we define the state value function at counted rewards, starting from the state s, and following the policy afterwards until termination: V=((s) Ea. Pl 1 a.) so s].
S as the expected sum of dis-
Model-based RL (MBRL). MBRL algorithms address the supervised learning problem of es- timating the dynamics of the environment P (and sometimes also the reward function f) from data collected when interacting with the real system. The model's loss function is typically the lig-likelihood C(D; P) = t Σnn log P(s |s; a;) or Mean Squared Error (MSE) for deter- ministic models. The learned model can subsequently be used for policy search under the MDP
2

---

Published as a conference paper at ICLR 2025
1.0
0.8
0.6
a
b
C
0..4
0.2
0..0
Figure 2: LLM can perceive time patterns. The LLM (Llama 3-8B) is fed with 3 time series presenting distinct patters.. (a) Rectangular pulse. (b) Rectangular signal with constant sub sub-parts. (c) The fthigh dimension HalfCheetah under an expert policy. Tokens belonging to constant slots (or peaks) attend to all the similar that precede them, focusing more their first occurrence.
M = (S, (o:7). 7). This MDP shares the state and action spaces S,A, reward function r, with the true environment M, but learns the transition probability from the dataset D.
Large Language Models (LLMs). Within the field of Natural Language Processing, Large Lan- guage Models (LLMs) have emerged as a powerful tool for understanding and generating human- like text. An LLM typically defined as neural network model, often based on the trans- former architecture (Vaswani et al., 2017), that is trained on a vast corpus of sequences, U where each sequence U, = (M1, U2,.... ,i....,,.................... consists of tokens Uj from vocabulary V. Decoder-only LLMs (Radford et al., 2019; Dubey & the Llama 3 team., 2024) typically encode an autoregressive distribution, where the probability of each token is condi- tioned only on the previous tokens in the sequence, expressed as po (U,) = I[ Pe(u,) [00-11). The
U...... UN) UU..U....
parameters 0 are learned by maximizing the probability of the entire dataset, po I (U.).. Every LLM has an associated tokenizer, which breaks an input string into sequence tokens, each belonging to V.
In-Context Learning (ICL). In or- der to use trajectories inputs in ICL, we use the tokenization of
Algorithm ICL, (Liu et al., 2024b; Gruver et al., 2023b)
Input: Time series (s:) LLM po, sub- socabuaayulary Vnum
1. Tokenize time series f. = ... .f......""
time series proposed in Gruver et al. (2023b) and Jin et al. (2024). This
2. logits Po(
approach uses subset of the LLM vocabulary Vnum representing digits tokenize the time series (Algo-
3. [P(Xi+ [E...... x0) wo)ses t- softmax(logits(Vmmm
Return: (P(Xi+ 12, ................................ xo) lis:
rithm 1). Specifically, given an univarite time series, we rescale it into a specific range (Liu et al., 2024b; Zekri al.; Requeima et al., 2024), encode it with digits, and concatenate each value to build the LLM prompt:
0.2513, 5.2387, 9.7889] + 1.5, 8.16, + 1150, 516, 850"
time series
rescaled
prompt
After the LLM forward pass, the logits corresponding to tokens in V can be used to predict categorical distribution over the next value as demonstrated in Liu et al. (2024c), thereby enabling uncertainty estimation.
3
ZERO-SHOT DYNAMICS LEARNING USING LARGE LANGUAGE MODELS
3.1 MOTIVATION
Before considering the multivariate trajectories of agents collected in RL environments, we first want to verify whether pre-trained LLM model is sensitive to the primitive univariate signals akin to those encountered in them. For this, we investigate the attention mechanism the Llama3 8B model (Dubey & the Llama 3 team., 2024) when feed it with different signals, including the

---

Published as a conference paper at ICLR 2025
periodic fthigh dimension from the HalfCheetah system (Brockman et al., 2016). By averaging the attention matrices over the 32 heads for each of the 32 layers of the multi-head attention in Llama3, observed distinct patterns that provide insight into the model's ''s focus and behavior (Fig. 2 shows selected attention layers for each signal). The attention matrices exhibit diagonal pattern, indicative of strong self-correlation among timestamps, and a subtriangular structure due to the causal masked attention in decoder-only transformers.
Further examination of the attention matrices reveals more intricate finding. Tokens within repeat- ing patterns (e.g., signal peaks, constant parts) not only attend to past tokens within the same cycle but also those from previous occurrences of the same pattern, demonstrating form of in-context learning. The ability to detect and exploit repeating patterns within such signals is especially valu- able in RI., where state transitions and action outcomes often exhibit cyclical or recurring dynamics, particularly in continuous control tasks. However, applying this insight to RI. presents two critical challenges related 1) the integration of actions into the forecasting process, and 2) handling the multivariate nature of RL problems. We now address these challenges by building on the insights from the analysis presented above.
3.2 PROBLEM SETUP
Given initial trajectory T = (so, do, 81, 01, .2......... (T1, 31 of length T, with 8; S, a, = (s;; ) E A', where the policy π is fixed for the whole trajectory, and r4 E R, we want to predict future transitions: given (sr-1, ay--1 predict the next state and reward (sr,rr) and subsequent transitions autoregressively. For simplicity we first omit the actions and the reward, focusing instead on the multivariate sequence n" (So. si,..., . ST) where
assume that the state dimensions independent. Later, we show how to relax the assumptions of omitting actions and rewards, as well state independence, which is crucial for applications in RL. The joint probability density function of T- can be written as:
P(T*) = (o(80) 11- PT (selstt-
(1)
where PT (sclst-1) = SaEA P(alss-)))(s(((sL..))))
Using the decoder-only nature of the in-context learner defined in Section 2, we can apply Algorithm to each dimension of the state vector to infer the transition rule of each visited state in n" conditioned on its relative history: for all j (1,....d.)
HalfCheetah
1.00
(009
651
0.50
0.25
(P (s/ s/
(2) 30.5
0.00
where are the fixed parameters of the LLM used as an in-context learner, and T its context length. Assum- ing complete observability of the MDP state, the Marko- vian property unveils an equivalence between the learned transition rules and the corresponding Markovian ones:
-0.55
Po(sy|se-1..... $1, sp) = (s
Figure 3: The covariance matrix from an expert dataset in the Halfcheetah en- vironment indicates linear correlations between state-action features.
This approach, that name vICL (for vanilla ICL), thus applics Algorithm on each dimension of the state indi- vidually, assuming their independence. Furthermore, the action information is integrated-out (as depicted in Eq. (1)), which in theory, limits the application scope of this method to quantities that only depend on policy through the expectation over actions (e.g., the value function (s)). We address these limitations in the next section.
On the zero-shot nature of DICL. Our use of the term "zero-shot" aligns with the literature on LLMs and time series (Gruver et al., 2023a), indicating that we do not perform any gradient up- dates or fine-tuning of the pretrained LLM's weights. Specifically, we adopt the dynamical sys- tems formulation of ICL as studied in Li et al. (2023), where the query consists of the trajectory " . s " and the label is the subsequent value s/-
In practice states and actions real valued vectors spanning space of dimensions respectively d, and d. S Rd R
4

---

Published as a conference paper at ICLR 2025
HalfCheetah
multi-step
bioot
fshin
fshin
bthigh
600
rooty
50
350
400
rootz
multi-step
400
bshin
fbot
200
athigh
bloot
Rhigh
0
50
350
400
0
(a) Multi-step error.
(b) Predicted trajectories.
(c) Time.
groundtruth
VICL
DICL-(s)
DICL-(s, a)
MLP
Figure 4: PCA-based DICL achieves smaller multi-step error in less computational time. We compare DICL-(s) and DICL-(s, a) using number of components equal to half the number of features, with the vanilla approach vICL and an MLP baseline. (Llama 3-8B).
3.3 STATE AND ACTION DIMENSION INTERDEPENDENCE
In this section we address the two limitations of vICL discussed Section 3.2 by introducing Dis- entangled In- In-Context Learning (DICL), a method that relaxes the assumption of state feature inde- pendence and reintroduces the action by employing strategies that aim to map the state-action vector to latent space where the features are independent. We can then apply vICL., which operates under the assumption of feature independence, to the latent representation. An added benefit of using such a latent space is that it can potentially reduce the dimensionality, leading to speed-up of the overall approach.
While sophisticated approaches like disentangled autoencoders could be considered for DICL, in this work we employ Principal Component Analysis (PCA). In fact, the absence of pre-trained mod- els for this type of representation learning requires training from scratch dataset. This goes against our goal of leveraging the pre-trained knowledge of LLMs and ICL. Instead, find that PCA, which generates new linearly uncorrelated features and can reduce di- mensionality, strikes a good balance between simplicity, tractability, and performance (Fig. 3 and Fig. 4). Nonetheless, DICL is agnostic this aspect and any transformation yp that disentangle features can be used in place of PCA. In the rest of the paper we present two variants of DICL:
potentially large
DICL-(s, a), which applies the rotation matrix of PCA to the feature space of states and actions and then runs Algorithm in the projection space of principal components;
DICL-(s), which applies the same transformation solely to the trajectory of states, This is useful in settings in which integrating the actions is not necessary, as when we only want to estimate the value function //8 (s).
3.4 AN ILLUSTRATIVE EXAMPLE
In this section, we aim to challenge our approach against the HalfCheetah system from the MuJoCo Gym environment suite (Brockman al., 2016; Todorov al., 2012). All our experiments conducted using the Llama series of models (Dubey & the Llama 3 team., 2024). Fig. 4a shows the average MSE a prediction horizon of h E (1,..., shows predicted trajectories for selected state dimensions of the HalfCheetah system (the details of the experiment, the metries and the remaining state dimensions deferred to Appendix F).
,,. 20} steps for each state dimension. Fig. 4b
We first observe that the LLM-based dynamics forecasters exhibit a burn-in phase (NU 70 steps in Fig. 4b) that is necessary for the LLM to gather enough context. For multi-step prediction, Fig, 4a, showing the average MSE over prediction horizons and trajectories, demonstrates that both versions of DICL improve over the vanilla approach and the MLP baseline trained on the context data, in almost all state dimensions. Indeed, we hypothesize that this improvement is especially brought by the projection in a linearly uncorrelated space that PCA enables. Furthermore, we also leveraged the
A detailed discussion of alternative approaches PCA is provided in Appendix C.

---

Published as a conference paper at ICLR 2025
dimensionalityy reduction feature by selecting number of components equal to half the number of the original features dg + dg (or dg in DICL-(s)). This results significant decrease in the computational time of the method without loss of performance, as showcased by Fig. 4c.
LLMs comparison. In Table we compare the perfor- mance obtained by the baselines and DICL when using different LLMs. Similarly to Fig. 4a, the scores are cal- culated as the average over given prediction horizon h across all dimensions (refer to Appendix F for details on the MSE, and Appendix G for details on the KS statistic).
Metrics
LLaMA
MSE/ 10 34
KS/10-2
vICL
3.2-1B
384 31
52 +
Note that similarly to Fig. 4, we use PCA-based dimen-
3,2-3B
399 + 40
54 #
sionality reduction for both DICL-(8, a) and (ICL-(s)) in
3,1-8B
380 + 32
53
this experiment, reducing the original number of features
3-8B
375 + 30
53 1
by half. Overall, we can see that DICL, especially the
3.1-70B
392 35
55
DICL-(8, a) version, demonstrates improved calibration compared to both vICL and the MLP baselines, thanks
DICL-(s)
3.2-1B
389 38
50 7
to the disentangling effect of PCA. Moreover, DICL-
3,2-3B
404 + 41
+1
(s) with the 3.1-70B model achieves the lowest Mean
3,1-8B
372 44
50
Squared Error (MSE) of 3.59. Nonetheless, DICL-(8, a)
3-8B
370 - 36
50
exhibits the highest MSE across all models. This is likely due the additional error introduced by predicting ac- tion information, thereby modeling both the dynamics and the data-generating policy. This aspect differs from the MLP baseline, which is provided with real actions at test time (acting as an oracle), and from DICL-(s) and vICL, which operate solely on states. We show the de- tailed results of this ablation study in Appendix H. Notice that we exclusively used LLMs based the LLaMA se- ries of models (Dubey the Llama team., 2024). This was strategic choice due the LLaMA tokenizer, which facilitates our framework by assigning a separate token to each number between 0 and 999. For other LLMs, al- gorithms have been suggested in the literature to extract transition rules from their output logits. For example, the Hierarchical Softmax algorithm (Liu et al., 2024b) could be employed for this purpose.
3.1-70B
359 33
54
DICL-(s,
3,2-1B
449
+ 37
46 5
3,2-3B
450
47
48 + 6
3.1-8B
412
+ 39
45 + 6
3-8B
418
46
46 +
3.1-70B
428
+ 47
47 +
baseline MLP
406 + 59
55 +
different LLMs. Results average over episodes from each one of 7 D4RL (Fu al., 2021) tasks. means lower the better. The best average score is shown bold. We show the average score + the 95% Gaussian confidence interval.
Table 1: Comparison
USE-CASES IN REINFORCEMENT LEARNING
As explored in the preceding sections, LLMs can be used as accurate dynamics learners for propri- oceptive control through in-context learning. We now state main contributions in terms of the integration of DICL into MBRL. First, we generalize the return bound of Model-Based Policy Op- general
timization (MBPO) (Janner et al., 2019) the analyze our method. Next, we leverage the LLM to augment the replay buffer of an off-policy RL al- gorithm, leading more sample-efficient algorithm. In second application, we apply our method to predict the reward signal, resulting in hybrid model-based policy evaluation technique. Finally, we show that the LLM provides calibrated uncertainty estimates and conclude with a discussion of our results.
of multiple branches and use it to
4.1 THEORETICAL ANALYSIS: RETURN BOUND UNDER MULTI-BRANCE ROLLOUTS
When using a dynamics model in MBRL, one ideally seeks monotonic improvement guarantees, ensuring that the optimal policy under the model is also optimal under the true dynamics, up to some bound. Such guarantees generally depend on system parameters (e.g., the discount factor 7), the prediction horizon k, and the model generalization error Em. As established in Janner et al. (2019) and Frauenknecht et al. (2024), the framework for deriving these thcoretical guarantees is the one of branched model-based rollouts.

---

Published as a conference paper at ICLR 2025
A branched rollout return s/brach(rri [r]] of policy is defined in Janner et al. (2019) as the return of a rollout which begins under the true dynamics P and at some point in time switches to rolling out under learned dynamics P for lo steps.
interested in studying a more general branching scheme that will be later used to analyze the results of our data-augmented off-policy algorithm. We begin by defining the multi- branch rollout return.
For our LLM-based dynamics learner,
Definition 4.1 (Multi-branch rollout return). The multi- branch rollout return 7/ "p. [T]] of policy π is defined as the expected return over rollouts with the following dynamies:
1pkr(T]] Blbranches]
1. for 1 < T, where T is the minimal context length, the rollout follows the true dynamics P.
Figure 5: Multi-branch return. The rollout following the true dynamics P is shown in blue. The branched roll- outs following LLM-based dynamics Pium are in purple. Branched rollouts can overlap. with the expectation over the overlapping branches as the return.
2. for t 2 T, with probability p. the rollout switches to the LLM-based dynamics Pimm for k steps. otherwise the rollout continues with the true dynamics P.
These different rollout realizations, referred to as branches, can overlap, meaning that multiple LLM-based dynamics can run in parallel if multiple branchings from the true dynamics occur within the k-step window (see Fig. 5).
With this definition, we now state our main theoretical result, consisting of a return bound between the true return and the multi-branch rollout return.
Theorem 4.2 (Multi-branch return bound). Let T be the minimal length of the in-context trajecto- ries, p E [0. 1] the probability that given state is branching point. We assume that the reward is bounded and that the expected total variation between the LLM based model and the true dynamics
under policy is bounded at each timestep by maxt2T Rg--p' "a>a Drv( P(-ss, a) [m a))] S eqm (T). Then under multi-branche rollout scheme with a branch length of k. the return
bounded as follows:
7" max p
In((T) ( [
(3)
1 1-7
where Fmax max r(s,
Theorem 4.2 generalizes the single-branch return presented in Janner et al. (2019), incorporating an additional factor of the prediction horizon i due to the presence of multiple branches, and directly accounting for the impact of the amount of LLM training data through the branching factor Addi- tionally, the bound is inversely proportional to the minimal context length T, both through the power in the discount factor 7² and the error term (T). Indeed, the term 6 (TT) corresponds to the generalization error of in-context learning. Several works in the literature studied it and showed that typically decreases in (T--1/2 with T the length of the context trajectories (Zekri et al., 2024; Zhang et al., 2023e; Li et al., 2023).
4.2 DATA-AUGMENTED OFF-POLICY REINFORCEMENT LEARNING
In this section, we show how DICL can be used for data augmentation in off off-policy model-free RL algorithms such as Soft Actor-Critic (SAC) (Haarnoja et al., 2018). The idea is to augment the replay buffer of the off-policy algorithm with transitions generated by DICL, using trajectories already collected by previous policies. The goal is improve sample-efficiency and accelerate the learning curve, particularly the early stages of learning as the LLM can generate accurate transitions from small trajectory. We name this application of our approach DICL-SAC.
As defined in Corrado & Hanna (2023), augeeneeeented oof-po RL involves perturbing previ- ously observed transitions to generate new transitions, without further interaction with the environ- ment. The generated transitions should idcally be diverse and feasible under the MDP dynamics to enhance sample efficiency while ensuring that the optimal policy remains learnable.
7

---

Published as a conference paper at ICLR 2025
1e3
Pendulum
1e3
HalfCheetah
1e3
Hopper
1.5
0.5
Return
Return
4
Return
1.0
-1.0
0.5
-1.5
0.0
0.0
0.5
0.25
0.75
1.00
0.25
0.50
0.75
Step
1e4
Step 1e5
Step
le5
DICL-SACIa 5%)
DICL-SAC(e= 10%)
DICL-SAC(o (aA4C(==25))
Data-augmented off-policy RL. In the early stages of training DICL-SAC improves the sample efficiency of SAC on three Gym control environments. Due to the intensive use of the LLM within DICL-SAC, conducted this experiment using the Llama 3.2- 1B model.
Figure 6:
Algorithm 2 (DICL-SAC) inte- grates multiple components to demonstrate novel proof-of- concept for improving the sample efficiency
Algorithm DICL-SAC
I: Inputs: LLM-based dynamics learner (e.g. DICL-(s),
batch size b, LLM data proportion a, minimal context length T, and maximal context length Tmax
SAC using DICL for data augmentation. Let T (So, Go, do, ro....,. ,T. a'eax'ss TTmax) be a real trajectory collected with a fixed policy Th. sampled from the real transitions being
2: Initialize policy nT. critic Qu- replay buffer R, and
LLM replay buffer Riim, and context size Tmax
for L...., -N-interactions do
4: New transition (st, at Tt, St+1) from T
Add (st, ae Tt: St+1) to
stored in a replay buffer R. We generate synthetic
6: Store auxiliary action ã; To((ss))
transitions
if Generate LLM data then
(3t, ae, Tt, 8 lrsssrsssmm
where
8:
Sample trajectory T (s0,... STans) from R {8i+1 JosisSs DICL-(() (7)
8441 is the next state generated by the LLM model applied on the trajectory
9:
10:
Add {(s;, a,, (11)) Rilm
of the states only, ā, is action
11:
end if
sampled from the data collection
12:
if update SAC then
policy na (-Iss), and T is the minimal
13:
Sample batch B of size b from R Sample batch Bum of size -b from Rum Update & and zi on B U Bum
context length. These transitions are then stored a separate replay buffer
14:
15:
Rum- At given update frequency,
16:
end if
DICL-SAC performs G gradient updates using data sampled from R
17:
end for
and a% gradient updates using data sampled from Rum Other hyperparameters our method include the LLM-based method (vICL, DICL-(s) DICL-(s, a)), how often generate new LLM data and the maximal context length Tiax (see Appendix D for the full list of hyperparameters)
Fig. compares the return curves obtained by DICL-SAC against SAC in three control environ- ments from the Gym library (Brockman et al., 2016). As anticipated with our data augmentation approach, observe that our algorithm improves the sample efficiency of SAC the beginning of training. This improvement is moderate but significant in the Pendulum and HalfCheetah envi- ronments, while the return curves tend be noisier in the Hopper environment. Furthermore, as the proportion of LLM data a increases, the performance of the algorithm decreases (particularly in HalfCheetah), as predicted by Theorem 4.2. Indeed, larger proportion of LLM data correlates with higher probability of branching p. more branching points will be sampled throughout the training. Regarding the other parameters our bound in Theorem 4.2, we set T 1, meaning L.MM-generated transitions are added to Rum. and k 1 to minimize LLM inference time.
4.3 POLICY EVALUATION
In this section show how DICL can be used for policy evaluation.
8

---

Published as a conference paper at ICLR 2025
System engineers often presented with several policies to test on their systems. On the one hand, off- off-policy evaluation (e.g., Uchara et al. (2022)) involves using historical data collected from a dif- ferent policy to estimate the performance of a target policy without disrupting the system. However, this approach is prone to issues such as distributional shift and high variance. On the other hand, on- line evaluation provides direct and unhiased comparison under real conditions. System engineers often prefer online evaluation for a set of pre-selected policies because it offers real-time feedback and ensures that deployment decisions are based on live data, closely reflecting the system's n true performance in production. However, online evaluations can be time-consuming and may temporar- ily impact system performance. To address this, we propose a hybrid approach using LLM dynamics predictions obtained through ICL to reduce the time required for online evaluation: the initial phase of policy evaluation is conducted as a standard online test, while the remainder is completed offline using the dynamics predictions enabled by the LLM's ICL capabilities.
Fig. 7 illustrates the relative error in value ob- tained by predicting the trajectory of rewards for k steps, given context length of 500. When k 500, complete the remaining steps of the 1000-step episode using the ac- tual rewards. For the two versions of DICL,, the reward vector is concatenated to the feature space prior applying PCA. In the Hopper en- vironment, it is evident that predicting the re- ward trajectory alone challenging task for the vanilla method vICL. On the contrary, both (s) and DICL-(s, a) effectively capture some of the dependencies of the reward signal on the states and actions, providing a more ro- bust method for policy evaluation, and match- ing the MLP baseline that has been trained on a dataset of transitions sampled from the same policy. However, in HalfCheetah we observe that the vanilla method largely improves upon both the baseline and DICL. We suspect that this is due to the fact that the reward signal is strongly correlated with the rootæ dimension in HalfCheetah, which proved to be harder to predict by our approach, as can be seen in Fig. 4a.
Hopper
HalfCheetah
40
20
9g
30
15
20
10
200
400
VICL DICL-(s)
DICL- (s,a)
MLP
Figure 7: Policy evaluation with DICL. Relative eгтor on the predicted value over k = 500 steps, with context length of T 500. This experiment is conducted using the Llama 3-8B model.
Note that the experimental setup that we follow here is closely related to the concept of Model-based Value Expansion (Feinberg et al., 2018; Buckman et al., 2018), where we use the dynamics model to improve the value estimates through an n-step expansion in an Actor Critic algorithm.
4.4 CALIBRATION OF THE LLM UNCERTAINTY ESTIMATES
An intriguing property observed in Fig. 4b is the confidence interval around the predictions. As detailed in Algorithm 1, one can extract a full probability distribution for the next prediction given the context, enabling uncertainty estimation in the L.M's predictions. Notably, this uncertainty is pronounced at the beginning when context is limited, around peaks, and in regions where the average prediction exhibits large errors. We explore this phenomenon further in the next section by evaluating the calibration of the LLM's uncertainty estimates.
Calibration is known be an important property of a dynamics model when used in reinforcement learning (Malik et al., 2019). In this section, we aim to investigate whether the uncertainty esti- mates derived from the LLM's logits are well-calibrated. We achieve this by evaluating the quantile calibration (Kuleshov et al., 201 8) of the probability distributions obtained for each L.-M-based method.
Quantile calibration. For regression problem with variable y y/ R. and model that outputs a cumulative distribution function (CDF) F; over (where i indexes data points), quantile calibration implies that yi (groundtruh)) should fall within p%6-confidence interval p% of the time:
EN S +p for all p [0 as N 00
(4)
N
9

---

Published as conference paper at ICLR 2025
where F [0, Y denotes the quantile function (p) inf p < (1) for all p E [0, 1]. and N the number of samples.
HalfCheetah bfoot
1.0
0.8
Proportion
LLMs are well-calibrated forecasters. Fig. 8 shows the reliability diagram for the bfoot dimension of the HalfCheetah system. The overall conclusion is that, regardless of the LLM-based sub-routine used predict the next state, the uncertainty estimates derived from the LLM's logits are well-calibrated of quantile calibration. Ideally, forecasters should align with the diagonal in Fig. 8, which the LLM approach nearly achieves. Further- more, when comparing with naive baseline (the details deferred to Appendix G), the LLM-forecaster matches the baseline when it's already calibrated, and improves over it when it's not. To quantify forecaster's calibration with point statistic, we compute the Kolmogorov-Smming goodness-of-ft Eq. (10), shown in the legend of Fig. 8.
propon
0.6
0.4
0.2
0.0
0.00
0 0.25
0.50
0.75
1.00
terms
quantile
Nave ks-
DICL-66
DICL(sa)
Figure 8: Quantile calibration reliability día- gram. The LLM (Ilama 8B) uncertainty esti- mates are well-calibrated. Vertical lines show the Kolmogorov Dolmogorov-Smirovv Smirnov statistic for each fit.
DISCUSSION
By introducing the DICL framework, our goal is to bridge the gap between MBRL and LLMs. Our study raises multiple open questions and future research directions. Notably, the choice of the feature transformation crucial for improving performance in specific applications. We plan explore transformations that capture not only linear but also non-linear dependencies, such as AutoEncoders, as discussed in Appendix C. Another possible direction is the integration of textual context information into the LLM prompt. This approach has been shown enhance the overall pipeline for time series forecasting (Jin et al., 2024; Xue & Salim, 2023) and policy learning (Wang al., 2023).
Besides this, our algorithm DICL-SAC performs data augmentation by applying the LLM to gen- erate next states Eq. (2). This operation requires a total of d, calls the LLM (or after the oo transformation) to generate Tmax T transitions, as the time steps can be batched. This approach assumes fixed policy in the context, allowing the L.LM to implicitly learn p using only the states. Looking ahead, future research direction is to explore how apply DICL to MBRL by replacing the dynamics model with an LLM. Naively applying DICL-(s, a) would require (Tmax - 'T' dd calls to the LLM, transitions need to be predicted sequentially when actions change. This results in an extremely computationally expensive method, making it infeasible for many applications. There- fore, further research is needed to make this approach computationally efficient.
CONCLUSION
In this paper, we ask how we can leverage the emerging capabilities of Large Language Models to benefit model-based reinforcement learning. We build on previous work that successfully conceptu- alized in in-context learning for univariate time series prediction, and provide a systematic methodol- ogy to apply ICL an MDP's dynamics learning problem. Our methodology. based on a projection of the data in a linearly uncorrelated representation space, proved to be efficient in capturing the dynamics of typical proprioceptive control environments, addition to being more computationally efficient through dimensionality reduction.
To derive practical applications of our findings, we tackled two RL use-cases: data-augmented off- policy RL, where algorithm DICL-SAC improves the sample efficiency of SAC, and benefits from a theoretical guarantee under the framework of model-based multi-branch rollouts. Our sec- ond application, consisted in predicting the trajectory of rewards in order perform hybrid online and model- moodel-based policy evaluation. Finally, we showed that the LLM- based dynamics model also provides well-calibrated uncertainty estimates.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
The authors extend their gratitude to Nicolas Boullé for insightful discussions on the initial con- cepts of this project, as well as to the authors of the paper (Liu et al., 2024c) (Toni J.B. Liu, Nicolas Boullé, Raphaël Sarfati, Christopher J. Earls) for providing access to their codebase. The authors also appreciate the anonymous reviewers and meta-reviewers for their valuable time and constructive feedback. This work was made possible thanks to open-source software, including Python (Van Rossum & Drake Jr, 1995), PyTorch (Paszke al., 2019), Scikit-learn (Pedregosa et al., I), and CleanRL (Huang et al., 2022).
REPRODUCIBILITY STATEMENT
In order to ensure reproducíbility we release the code https::/githu..com//aaaa implementation details and hyperparameters are listed in Appendix D.
http://hhbbb.oombbbbeecccchhaa.. The
REFERENCES
Abdelrahman Abdelhamed, Mahmoud Afifi, and Alec Go. What do you see? enhancing zero-shot image classification with multimodal large language models. arXiv preprint arXiv:2405. 15668, 2024.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? Investigations with linear models, May 2023. URL http: / //aariivv www.v..cr//aa........oo abs 221..16611......1111111111111111 15661. arXiv:2211.15661 [cs].
Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. In International Con-
ference Learning Representations, 2021.
Abdelhakim Benechehab, Albert Thomas, and Balázs Kégl. Deep autoregressive density
nets vs neural ensembles for model-based offline reinforcement learning. arXiv preprint arXiv:2402.02858, 2024.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAl gym, 2016.
Tom B. Brown, Benjamin Mann, Nick Ryder. Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners, July 2020. URL httt:://arkhe oro.orrggae abs waabs/200 2005.0..00 4165. arXiv:2005.144 14165 [cs].
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, pp- 8234 8244, Red Hook, NY, USA, 2018. Curran Associates Inc
Yuji Cao, Huan Zhao, Yuheng Cheng. Ting Shu, Yue Chen, Guolong Liu, Gaoqi Liang, Junhua Zhao, Jinyue Yan, and Yun Li. Survey on large language model-enhanced reinforcement learning: Concept, taxonomy, and methods, 2024. URL httpt:tpp///oc http:p/s/////....e.eeeemo
rxiv. wwwpwiiwwog//bb//eeeeeo 2404 00282.
Thomas Carta, Clément Romac, Thomas Wolf, Sylvain Lamprier, Olivier Sigaud, and Pierre-Yvee Oudeyer. Grounding large language models in interactive innernattiveeniioommentt envionmmnns with online reinforcement learning, 2023. URL httpt:::/// / httpvviivvvv.///. org/ abs 02662.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models, 2022. URL httph:tt////ccccct http///axivvvoo
/tt/2222222 09..0.094488 09481
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, 2021. URL. https: teattps://aate http://wwwwvw..oa//////////////s 200.e/..5.......5..5..

---

Published as a conference paper at ICLR 2025
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems 31, pp.
Curran Associates, Inc., 2018.
Julian Coda-Forno, Marcel Binz. Zeynep Akata, Matthew Botvinick, Jane X. Wang, and Eric Schulz. Meta-in-context learning in large language models, May 2023. URL http:://aaaat. /ttp/i.iiii.////////.........//. abs/2305 12907. arXiv:2305. 12907 [cs].
Nicholas E Corrado and Josiah P Hanna. Understanding when dynamics-invariane data augmenta- tions benefit model-free reinforcement learning updates. arXiv preprint arXiv:2310.17786, 2023.
Abhimanyu Das, Weihao Kong, Rajat Sen, and Yichen Zhou. A decoder-only foundation model for time-series forecasting, April 2024. URL arXiv:2310. 10688 [es].
Mare Peter Deisenroth and Carl Edward Rasmussen. PILCO: A model-based and data-efficient approach policy search. In Proceedings of the International Conference Machine Learning, 2011.
Andreas Draeger, Sebastian Engell, and Horst Ranke. Model predictive control using neural net- works. IEEE Control Systems, 15:61-66, 1995. ISSN 1066033X. doi: 10.1109/37.46621..
Abhimanyu Dubey and the Llama 3 team. The llama 3 herd of models, 2024. URL https: //arxiv.org/aaaaaa.ie..a aab/////rgga////////ooo/... 2400 21783.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I. Jordan, Joseph E. Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning, 2018.
URL https::tp///mm /ttp/aaxxivv.ro/ee/aa org/ abs/1803. 00101.
Xueyang Feng, Zhi- Yuan Chen, Yujia Qin, Yankai Lin, Xu Chen, Zhiyuan Liu, and Ji-Rong Wen. Large Language Model-based Auean-Agent Collaboration for Complex Task Solving, February 2024. URL http: /arxiv.viiv....aa/aaaaaaaaaaa.va org aacp//ba//// /2002. 12914. arXiv:2402.12914 [cs].
Bernd Frauenknecht, Artur Fisele, Devdutt Subhasish, Friedrich Solowjow, and Sebastian Trimpe. Trust the model where trusts itself - model-based actor-critic with uncertainty-awar rollout adaption, 2024. URL https: wtw.aaiiiiiiv./.im http://aaa.i....aa////aaaaa////a 19014.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2021. URL https: /ttp::::////eeeeeeewee net forum?id-
px0- N000-N3.00 N3_ _j.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2062- PMLR, 09-15 Jun 2019. URL https: htttpps:/ppocceedngg mmlr.
press/ /elas/v97/0 fu htt///fjmmoooooooooooo html.
Yarin Gal, Rowan McAllister, and Carl Edward Rasmussen. Improving PILCO with Bayesian neural network dynamics models. In Data- Eifa--fffiieent Machine Learning workshop. International Confer- ence on Machine Learning, 2016.
Shivam Garg, Dimitris Tsipras. Percy Liang, and Gregory Valiant. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes, August 2023. URL http://arxiv. org//bs/// hat//bb/220000 /208..22... 01 01666. #xXiv:2208.01066 [cs].
Panagiotis Giadikiaroglou, Maria Lymperaiou, Giorgos Filandrianos, and Giorgos Stamou, Puzzle solving using reasoning of large language models: A survey. arXiv preprint arXiv:240..11291, 2024.
Nate Gruver, Mare Finzi, Shikai Qiu, and Andrew Gordon Wilson. Large Language Models Are Zero-Shot Time Series Forecasters, October 2023a. URL http: htttt:t//aaiiiii wtwp:///..gg/////////b//b/se 07820. arXiv:2310.0788 [cs].
12

---

Published as a conference paper at ICLR 2025
Nate Gruver, Mare Anton Finzi, Shikai Qiu, and Andrew Gordon Wilson. Large language models zero-shot time series forecasters, In Thirty-seventh Conference on Neural Information Processing Systems, 2023b. URL https:///penteviee net/ Eorum weweruummimmmddaee 2et.::dmmmmmmmmmmmmm
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Biance and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 2450-2462. Curran Associates, Inc., 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with Stochastic Actor. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870. PMLR, 10-15 Jul 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th In- ternational Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research. pp.
2019.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021. URL
https:/ / /openre// ew net/ forum? Forum?id=0aataia d=0oabwy7b0u..
Stefan Hegselmann, Alejandro Buendia, Hunter Lang, Monica Agrawal, Xiaoyi Jiang, and David Sontag. TabLLM: Few-shot Classification of Tabular Data with Large Language Models, March 2023. URL wwttttptttaaaae
orgg/bss//////ooooooooooooooooes /221 10723. arXiv:2210.10723 [cs].
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017. URL https::/
net/ for.mmdd=Syy22useggee
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Ki- nal Mehta, and João G.M. Araújo. Cleanrl: High-quality single-file implementations of deep reinforcement learning algorithms. Journal of Machine Learning Research, 23(274):1-18. 2022. URL /t org hatgg//ppers////e /vteee///v///// /t////2231/21/// 22/2242112222 html.
Louis Martin Hugo Touvron and the Llama 2 team. Llama Open foundation and fine-tuned chat models, 2023.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. In H. Wallach, H. Larochelle. A. Beygelzimer, F. d'Alchéé Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Offline Reinforcement Learning as One Big Sequence Modeling Problem, November 2021. URL 02039. arXiv:2106.02039 [cs].
Ming Jin, Shiyu Wang. Lintao Ma, Zhixuan Chu, James Y. Zhang, Xiaoming Shi, Pin-Yu Chen, Yuxuan Liang, Yuan-Fang Li, Shirui Pan, and Qingsong Wen. Time-llm: Time series forecasting by reprogramming large language models, 2024.
Shyam Sundar Kannan, Vishnunandan L. N. Venkatesh, and Byung- Cheol Min. SMART-LLM: Smart Multi-Agent Robot Task Planning using Large Language Models, March 2024. URL http: /awiiv.ivvvv..////v/ org/ abs/ /0ls/23300 22099 10062. arXiv:2309. 1.0220.00000......... [cs].
Balázs Kégl, Gabriel Hurtado, and Albert Thomas. Model-based micro-data reinforcement learn- ing: what are the crucial model properties and which model to choose? In International Confer- ence on Learning Representations, 2021. URL https:/ http////ppe///e/e/e////e/// net forum?ide
P5cy1694568.
13

---

Published as a conference paper at ICLR 2025
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Bal- can, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp- 21810-21823. Curran Associates, Inc., 2020. URL https: http:////oooeeeeeeeeeeeeee neurips. scc/paper/2020) (t/////fil//// 400
Hyunjik Kim and Andriy Mnih. Disentangling by factorising, 2019. URL https: //rxiv.. org/abs// aba/ 110..800..000 05983.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua
Bengio and Yann LeCun (eds.), ), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. URL http: //arxii.org////// abs / 1412. 6980.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2022. URL https: www.xiiiorrg/as//./. a/1/112../1111..//../ 1112. 6114.
Volodymyr Kuleshov, Nathan Fenner, and Stefano Ermon. Accurate uncertainties for deep learning using calibrated regression. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume Proceedings of Machine Learning Research, pp. 2796-2804 PMLR, 10-15 Jul 2018. URL
mmlr.
press/ www./////kkllleeleeeeee html.
Minae Kwon, Sang Michael Xie, Kalesha Bullard, and Dorsa Sadigh. Reward design with language models, 2023. URL wwtptttt//// arxiv. org/ abtp:/bb//20000 00001.
Byung-Jun Lee, Jongmin Lee, and Kee-Eung Kim. Representation balancing offline model-based reinforcement learning. In International Conference on Learning Representations, 2021. URL
www.mmiddogpa2nrraa... net/ wttp://wwww.eeeeoeeteeeeeeeeeeee
Sergey Levine and Vladlen Koltun. Guided policy search. In Sanjoy Dasgupta and David McAllester (eds.), Proceedings of the 30th International Conference Machine Learming, volume 28 of Proceedings of Machine Learning Research, pp. 1-9, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL httph:t/p///.m /tt.////Ltttt////oeeeeeeee pttpppeeceeengggeeeeeeee www.m.r..rree.. www.spresssee..ee. wewrees//28/leeiine .hmml.
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Reseurch, pp. 19565-19594. PMLR, 23-29 Jul 2023. URL https: http::///occeeiieeee www.m.rresss....... http://22//1/22 .html.
Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code as Policies: Language Model Programs for Embodied Control, May 2023.
URL http: /tt///axiivv org/ abws/209..... 07753. arXiv:2209.07553 [es].
Jessy Lin, Yuqing Du, Olivia Watkins, Danijar Hafner, Pieter Abbeel, Dan Klein, and Anea Dragan. Learning to Model the World with Language, May 2024. URL http: org/ abs/ 2308. 01399. arXiv:2308.00 033999999999 [es].
Ruizhen Liu, Zhicong Chen, and Dazhi Zhong. Dromo: Distributionally robust offline model-based policy optimization. 2021.
Shaoteng Liu, Haoqi Yuan, Minda Hu, Yanwei Li, Yukang Chen, Shu Liu, Zongqing Lu, and Jiaya Jia. RL-GPT: Integrating Reinforcement Learning and Code-as-policy, February 2024a. URL
http: //arxivaavvva////////a/eaaaaaaaa wwt/oorg//////eee//eo abs//402.2222/ 19299. arXiv:2402.1922 [es].
Toni B. Liu, Nicolas Boullé, Raphaël Sarfati, and Christopher Earls. LLMS learn governing principles of dynamical systems, revealing an in-context neural scaling law, 2024b.
Toni J. B. Liu, Nicolas Boullé, Raphaël Sarfati, and Christopher J. Earls. LLMs learn governing principles of dynamical systems, revealing an in-context neural scaling law, February 2024c. URL http::t///// http///avvvv..///e//.. wwwlvv.orrg/aaae hllg/abs/22e 00795. 40Xiv:2402.00795 [cs].
14

---

Published as a conference paper at ICLR 2025
Zuxin Liu, Jesse Zhang, Kavosh Asadi, Yao Liu, Ding Zhao, Shoham Sabach, and Rasool Fakoor. TAIL: Task-specific Adapters for Imitation Learning with Large Pretrained Models, October 2023. URL httttp://axxx..orgpg.. wwtp///////////oo///eeo///eeaee/ http//b//22133 /b35023..05005 05905. arXiv:2310.05905 [cs].
Runyu Ma, Jelle Luijkx, Zlatan Ajanovie, and Jens Kober. ExploRLLM: Guiding Exploration in Reinforcement Learning with Large Language Models, March 2024. URL http: //axxiv.. org/abs/2400 09583. arXiv:2403.0955 .05583 [es].
Ali Malik, Volodymyr Kuleshov, Jiaming Song, Danny Nemer, Harlan Seymour, and Stefano Er- mon. Calibrated Model-Based Deep Reinforcement Learning. In Kamalika Chaudhuri and Rus- lan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learn- ing. volume 97 of Proceedings of Machine Learning Research, pp. 4314-4323, PMLR, 09-15 Jun 2019. URL https: http://pccceeennseeeeeeeee mlr wwt.s.ml..pp..r.......... wttpppppsssss/asaaapppeepp htt.lll............
Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment- efficient reinforcement learning via model-based offline optimization. In International Confer- ence Learning Representations, 2021. URL https: http//ppeeeeeiiieeeeee/eeeeeeeet net wttp://wwwee//oumiiiie ffrummiddd
3hGNqpI4MM 4WS..
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are Sample-Efficient World Mod- els. September 2022. URL https: //ppnevvieeeeeeeeeeee net Fow.mmidd//uuuaabbbbb
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Pedregosa, G. Varoquaux, Gramfort, V. Michel, Thirion, O. Grisel, Blondel, Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830. 2011.
Eduardo Pignatelli, Johan Ferret, and Tim Rocktaschel. Assessing the Zero-Shot Capabilities of LLMs for Action Evaluation RL.
Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, Jan Achterhold, Joerg Stueekler, Michal Rolinek, and Georg Martius. Sample-efficient cross-entropy method for real-time planning. In Conference on Robot Learning 2020, 2020. URL https: http:://ccorccoce/ thub. 1o/
cor12020 sar12020/paper___ 20t-0000ppperr27/.. 217/.
Rudra P. K. Poudel, Harit Pandya, Chao Zhang, and Roberto Cipolla. guage Grounded World Model, November 2023. URL 17593v1.
LanGWM: Lan- halggg/bb///e/
Nooshin Pourkamali and Shler Ebrahim Sharifi. Machine translation with large language models: Prompt engineering for persian, english, and russian directions. arXiv preprint arXiv:2401.0842, 2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Machel Reid, Yutaro Yamada, and Shixiang Shane Gu. Can Wikipedia Help Offline Reinforcement Learning?, July 2022. URL http:/ http:///wwwwa..o//////////..10./ 12122 arXiv:2201.12122 [cs].
James Requeima, John Bronskill, Dami Choi, Richard E. Turner, and David Duvenaud. LLM Pro- cesses: Numerical Predictive Distributions Conditioned on Natural Language, May 2024. URL http:://axiis wrw.oorg/aaao.. httgg//b////// /205 12856. arXiv:2405. 12856 [es, stat].
15

---

Published as a conference paper at ICLR 2025
Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge you pack into the parameters of a language model? In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu (eds.). Proceedings of the 2020 Conference Empirical Methods in Natural Language Process- ing (EMNLP), P, pp. 1416-,426, Online, November 2020. Association for Computational Linguis- tics. doi: 10. 00.18663/1/2220.emmml emr1p-mm ma ii..ini.n 437.
URL ht.ps: //acl aetllanthhooogg org/ /2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms, 2017. https:/:/// http//axiiv/////////v///.//aaaaa org
Ruizhe Shi, Yuyao Liu, Yanjie Ze, Simon S. Du, and Huazhe Xu. Unleashing the Power of Pre- trained Language Models for Offline Reinforcement Learning, November 2023, URL http:
//arxiv.org// aww/aaaa..........a./..aaaaaataa 20587. 7. arXiv:2310.20587 [cs].
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM SIGART Bulletin, 2:160-163, 7 1991. ISSN 0163-5719. doi: 10.1145/122344412237, URL
bittps://d..ee org. / /10. www.111111114..... 122377.
Richard S Sutton, Csaba Szepesvári, Alborz Geramifard, and Michael Bowling, Dyna-style planning with linear function approximation and prioritized sweeping. Moore and Atkeson, 1992.
Albert Thomas, Abdelhakim Benechehab, Giuseppe Paolo, and Balázs Kégl. Fair model-based reinforcement learning comparisons with explicit and consistent update frequency. The Third Blogpost Track at ICLR 2024, 2024. URL https:/ / httpppppnwwwweeee.eeeop foret//oommmaa ShPNDZYND6.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 JEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, 2012. doi: 10.1100//ROS..012..3886..
Masatoshi Uehara, Chengchun Shi, and Nathan Kallus. A review of off-policy evaluation in rein- forcement learning, 2022. URL https: /ttpss///arxivii org at////b///// /ta.s//212... 06355.
Robert Vacarcanu, Vlad-Andrei Negru, Vasile Suciu, and Mihai Surdeanu. From Words to Numbers: Your Large Language Model Secretly A Capable Regressor When Given In-Context Examples, September 2024. URL http: //trx//aaiv/// org/ abs 2404 07544. arXiv:2404.07544 [es].
Guido Van Rossum and Fred L Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st Inter- national Conference on Neural Information Processing Systems, NIPS' 17, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Johannes Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent, May 2023. URL http: // arxiv. org/
07677. arXiv:2212.07677 [es].
Yen-Jen Wang, Bike Zhang, Jianyu Chen, and Koushil Sreenath. Prompt a Robot to Walk with Large Language Models, November 2023. URL http:://axixi org.vvorrg/aaa http://bb///000000 arXiv:2309.09961 [cs, eess].
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gug- ger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38 45, Online, October 2020. As- sociation for Computational Linguistics. URL https: wwww. aclweb. org/ holoogy 202 1200. emn 1 p-ddmos... 6.
16

---

Published as a conference paper at ICLR 2025
Yue Wu, Yewen Fan, Paul Pu Liang, Amos Azaria, Yuanzhi Li, and Tom M. Mitchell. Read and reap the rewards: Learning to play atari with the help of instruction manuals, 2024. URL https: / //t.axiiwrgg//aaaaaaa//.. als///b////// 22.02232.0000 04449.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An Explanation of In-context Learning as Implicit Bayesian Inference, July 2022. URL htttt://aaxiiii wrtpa///ggg/ab///1//// 02080. arXiv:21 11. 1,00000 [cs].
Hao Xue and Flora D. Salim. PrompiCast: A New Prompt-based Learning Paradigm for Time Series Forecasting, December 2023, URL http:://tt/ti../////////a/a// wtt.o/////./..oooo//r// 08964. arXiv:221000896 [es, math, stat].
Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foundation Models for Decision Making: Problems, Methods, and Opportunities, March 2023. URLhttpa
/ //a/aaaa //arxiiv.org abs 2303.04123 arXiv:2303.04129 [cs].
Yu Yang and Pan Xu. Pre-trained Language Models Improve the Few-shot Prompt Ability of Decision Transformer, August 2024. URL http:://rxlx .rrg abw//////////////////// 01402.
arXiv:2408.01402 [es].
Tianhe Yu, Garrett Thomas, Lantao Yu. Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 14129-14142. Curran Asso- neurips cc/papper/220/filee Pappe
ciates, Inc., 2020. URL pettps://proceeeeee a
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 28954-28967. Curran Asso- neurips ctcpppsccccpppee cc/ppaper//22///1leee
ciates, Inc., 2021. URL
#29a17977 11211101331113722233553535333333
Oussama Zekri, Abdelhakim Benechehab, and levgen Redko. Can Iims predict the convergence of stochastic gradient descent? In ICML 2024 Workshop on In-Context Learning.
Oussama Zekri, Ambroise Odonnat, Abdelhakim Benechehab, Linus Bleistein, Nicolas Boullé, and levgen Redko. Large language models as markov chains. arXiv preprint arXiv:2410.02724, 2024.
Xianyuan Zhan, Xiangyu Zhu, and Haoran Xu. Model-based offline planning with trajectory prun- ing. 2021.
Shenao Zhang, Sirui Zheng, Shuqi Ke, Zhihan Liu, Wanxin Jin, Jianbo Yuan, Yingxiang Yang, Hongxia Yang, and Zhaoran Wang. How Can LLM Guide RL. A Value-Based Approach, Febru- ary 2024.
www.ri...o.r/aa.ooooooo /2402 61 818.. arXiv:2402.161 16181 81 [cs].
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning, October 2023a. URL https:
// //taxxiv.rrg////eeelle/////ooltt abs. 2310.00 0961
Wenxuan Zhang, Yue Deng, Bing Liu, Sinno Jialin Pan, and Lidong Bing, Sentiment analysis in the era large language models: A reality check, 2023b. URL 2305. 15005.
abs/
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization, 2023c. URL https: //arxii///xvvv..// org/ abs 2205. 19420.
Zhaoheng Zheng, Jingmin Wei, Xuefeng Hu, Haidong Zhu, and Ram Nevatia. Large language mod- els are good prompt learners for low-shot image classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 28453-28462, 2024.
17

---

Published as a conference paper at ICLR 2025
Appendix
Outline. In Appendix A. we prove our main theoretical result (Theorem 4.2). We provide an extended related work in Appendix B. Additional materials about the state and action dimensions interdependence are given Appendix C. The implementation details and hyperparameters of our methods are given in Appendix D. Finally, we provide additional experiments about multi-step errors (Appendix F), calibration (Appendix G), the impact of the data collecting policy on the prediction emror (Appendix E), and details about the ablation study the choice of the LLM (Appendix H).
TABLE OF CONTENTS
Theoretical analysis
19
A.1 Proof of Theorem 4.2
19
B Related Work
21
State and action dimensions interdependence - additional materials
22
C.I Principal Component Analysis (PCA)
22
C.2 Independent Component Analysis (ICA)
C33 AutoEnnoder-based approach
23
C.4 Sensitivity Analysis
24
D Algorithms
24
D.1 Soft-Actor Critic
24
D.2 DICL-SAC
25
What is the impact the policy on the prediction error?
26
Multi-step prediction errors
26
G Calibration
27
H On the choice of the LLM
28
18

---

Published as a conference paper at ICLR 2025
A THEORETICAL ANALYSIS
A.1 PROOF OF THEOREM 4.2
We start by formally defining the LLM multi-branch return Mp.K.T- To do so, first denote A, the random event of starting k-step LLM branch at timestep and we denote X₃ the associated indicator random variable X, = [[A,]- We assume that the (X))e2T are independent. We then define the random event A that at least one of the & preceding timesteps has been branched, meaning that the given timestep i. belongs at least one LLM branch among the possible branches: A
U At- The LLM multi-branch return can then be written as follows:
T
(p. =
7 Ee Pt aart (r(s,,
t=0
Burnin phase gather minimal context T
E Mpx
[[A
Σ X X si aron [r(s adl
(5)
E X X1-1 i=1
l=T
average roward among the branches spanning timestep
[[A aur [r( anl
When no branch is spanning timestop
where p = = P(.ppt=1) with pa p4o the initial state distribution and P Pim((11 Pt
Before continuing, we first need to establish the following lemma.
Lemma A.1. (Multi-step Error Bound, Lemma B.2 in Frauenknecht et al. (2024) and Janner et al. (2019).) Let P and P be two transition functions. Define the multi-step error at time step starting from any initial state distribution Hl, as:
E Drv( (P ((Po))
with p0 = p0 Ho-
Let the one-step error time step 2 be defined as:
&r = Eg 1(1po)) [
and Eo = O.
Then, the multi-step error satisfies the following bound:
C1
E
2=
Proof. Let / > 0. We start with the definition of the total variation distance:
E, (Plo
/
(ss('(po)) ( ds'
JEES
/ (t-(ss)(-1(sspp))
ds ds'
JEES
J
((s'ls) (s-(() P (t ds ds'
<
JEES JES
(01)s)(-1(((1aa) pr ds ds
SES SES
19

---

Published as a conference paper at ICLR 2025
I Ivves sES ((ss))p()(((() P(s'(s) (s) p (s|no)
+P(s'(s) (s)po) P(s((s))(t11/ (s| ds
1
pt (s/mo) P(s'(s) - P(s'ls) ds ds'
EES AES
P(s'(s) (s - pr (slpo)) ds ds'
SES SES
1
-P((sss ds' p ds
SES 2 SES
(1 P(dl))) pp- (s|uo) ( (8)ma) ds [iv(((()o)) + Div((t-'((lpu)) ((1Ha)]|P(_1 (-)m0())m)))
SES
(
= +
Given that Eo Eo = 0, by induction we have:
E1 S,-
We now restate and prove Theorem 4.2:
Theorem A.2 (Multi- brati-bbrnnchhc return bound). Let T be the minimal length of the in-context trajecto- ries. E [0 the probability that a given state is branching point. We assume that the reward is bounded and that the expected total variation between the M-bsed model and the true dynamics under policy is bounded each timestep by maxe2T Bg-Pi,a>> [Drv( (([-ls, a) Puml--s. (s, a))] S eum (T)) Then under multi-branched rollout scheme with a branch length of k, the return is bounded as follows:
7 Tmaak?? p (T)
In(==)
Ip.. $2
Hm (T)) 1 1-9
where Tmax max a(s
Proof. Step 1: Expressing the bound terms of horizon-dependent errors.
xc g' BBa [r)s
(m](z)
_l .pp.k.
(T)
Σ
t=T
-Ex.
1(A
X, X1_ Pi Beri
Pi ae (r(s,, a:)
k Xx-
[r(se,
Ex.
[r + 1[A In (n
-Ex 1(4 Σ X 1=1
X2- se Pere ((88))a))
- (
ar)
20

---

Published as a conference paper at ICLR 2025
[[481
Ex (00
E CawPPtannn (r(0)) 0r)
Σ Xt-
Be the arnt [r(e,s at))
Σ Xiii 1=1
1|401
Σ Xi-i
Ei Xe-i i=1
T=T
Be (rl, [i)] - B Pin [r ail]
We then expand the integrals in the terms Eg Ee express it in terms of horizon-dependent multi-step model errors:
[r(s,, a.)l
E Bf mm [r(s,, ae)l and
Bac-Pi [r)) - E Pe amm [r)
-sss I r(s.a)(P'(s (8. (s.m(saa) da ds
aEd
Tmax [ JSES JEA
(P'(s.a) - -P, (s. a)) da ds
(P'(s) - (P))) (t)) da ds
S Tmax
JES DEA
JOES (P
Pr: (s)) ds
Tmax
$ 2rmax Drv(PL|| p P(
Step Simplifying the bound.
By applying Lemma A. we can bound the multi-step errors using the bound on one-step errors:
Drv(P')| f.0m) < i i (T) I & Eym (T)
(8)
Therefore, the bound becomes:
(n(=)
Jim k ( S 2 max mm (T)
Bx
1[A
Σ X
>
X
2r h Em (
[1[4]]]
t=
< 2am to Ell (T)
kp
t=
2 T
1 Tmack? 22 E (T)
B RELATED WORK
Model-based reinforcement learning (MBRL). MBRL has been effectively used in iterated batch RL by alternating between model learning and planning (Deisenroth & Rasmussen, 2011; Hafner al., 2021; Gal et al., 2016; Levine & Koltun, 2013; Chua et al., 2018: 18;; Janner et al., 2019;
21

---

Published as a conference paper at ICLR 2025
Kégl et al., 2021), and in the offline (pure batch) RL where we do step of model learning followed by policy learning (Yu et al., 2020; Kidambi et al., 2020; Lee et al., 2021; Argenson & Dulac-Arold, 2021; Zhan et al., 2021; Yu et al., 2021; Liu et al., 2021; Benechehab et al., 2024). Planning is used either at decision time via model-predictive control (MPC) (Draeger et al., 1995; Chua al., 2018; Hafner al., 2019; Pinneri et al., 2020; Kégl al., 2021), in the background where a model-free agent is learned on imagined model rollouts (Dyna: Janner et al. (2019); Sutton (1991); Sutton et al. (1992); Ha & Schmidhuber (2018)), or both. For example, model-based policy optimization (MBPO) (Janner et al., 2019) trains an ensemble of feed-forward models and generates imaginary rollouts to train a soft actor-critic agent.
LLMs in RL. LLMs have been integrated into reinforcement learning (RL) (Cao et al., 2024; Yang et al., 2023), playing key roles in enhancing drcision-making (Kannan et al., 2024; Pignatelli et al.;; Zhang et al., 2024; Feng et al., 2024), reward design (Kwon et al., 2023; Wu et al., 2024; Carta et al., 2023; Liu et al., 2023), and information processing (Poudel et al., 2023; Lin et al., 2024). The use of LLMs as world models is particularly relevant to our work. More generally, the Transformer architecture (Vaswani et al., 2017) has been used in offline RL (Decision Transformer Chen et al. (2021); Trajectory Transformer Janner et al. (2021)). Pre-trained LLMs have been used to initialize decision transformers and fine-tune them for offline RL tasks (Shi et al., 2023; Reid et al., 2022; Yang & Xu, 2024). As world models, Dreamer-like architectures based on Transformers have been proposed (Micheli et al., 2022: Zhang et al., 2023a; Chen et al., 2022), demonstrating efficiency for long-memory tasks such as Atari games. In text-based environments, LLMs have found multiple applications (Lin et al., 2024; Feng et al., 2024; Zhang et al., 2024; Ma et al., 2024), including using code- gode-gennratiig LLMs to generate policies in a zero-shot fashion (Liang et al., 2023; Liu et al., 2024a).
The closest work to ours is Wang et al. (2023), where system prompt consisting of multiple pieces of information about the control environment (e.g.. description of the state and action spaces, nature of the controller, historical observations, and actions) is fed to the LLM. Unlike our approach, which focuses on predicting the dynamics of RL environments, Wang et al. (2023) aim to directly learn a low-level control policy from the LLM. incorporating extra information in the prompt. Furthermore, Wang et al. (2023) found that only GPT-4 was usable within their framework, while we provide a proof-of-concept using smaller open LLMs such as Llama 3.2 IB.
ICL on Numerical Data. In-context learning for regression tasks has been theoretically analyzed in several works, providing insights based on the Transformer architecture (Li et al., 2023; von Os- wald et al., 2023; Akyürek al., 2023; Garg al., 2023; Xie al., 2022). Regarding time series forecasting, LLMTime (Gruver et al., 2023a) successfully leverages ICL for zero-shot extrapolation of oie-dimensional time series data. Similarly, Das al. (2024) introduce a foundational model for one-dimensional zero-shot time series forecasting, while Xue & Salim (2023) combine numerical data and text in a question-anse aanwerr format. ICL can also be used approximate a continuous density from the LLM logits. For example, Liu et al. (2024c) develop a Hierarchical softmax algorithm to infer the transition rules of uni-dimensional Markovian dynamical systems. Building on this work, Zekri et al. provide an application that predicts the parameter value trajectories in the Stochastic Gradient Descent algorithm. More relevant to our work, Requeima et al. (2024) presented LLMPro- cesses. method aimed extracting multi-dimensional distributions from LLMs. Other practical applications of ICL on numerical data include few-shot classification on tabular data (Hegselmann et al., 2023), regression (Vacareanu et al., 2024), and meta ICL (Coda-Forno et al., 2023).
C STATE AND ACTION DIMENSIONS INTERDEPENDENCE ADDITIONAL MATERIALS
C.I PRINCIPAL COMPONENT ANALYSIS (PCA)
Principal Component Analysis. PCA is a dimensionality reduction technique that transforms the original variables into a new set of variables, the principal components, which are linearly be ordered such that the first few retain most of the variation present in all of the original variables. Formally, given a data matrix X with n observations and p variables, PCA diagonalizes the covariance matrix C x'X'X to find
uncorrelated. The principal components
22

---

Published as a conference paper at ICLR 2025
the eigenvectors, which represent the directions of the principal components: PCA: - Z XW, where W are the eigenvectors of C. In our case, the data represents a dataset of states and actions given data collecting policy Th, while the p variables represent the state (eventually also the action) dimensions.
Ablation on the number of components. Fig. shows an ablation study the number of components used in the DICL-(s, a) method. Surprisingly. we observe a sharp decline in the average multi-step error (see Appendix F for a detailed definition) given only 4 components among 23 in the HalfCheetah system. The error then slightly increases for an intermediate number of components, be-
HalfCheetah
0.7
avg_ eererrr
0.6
fore going down again when the full variance recov- ered. This finding strengthens the position of PCA as our Disentangling algorithm of choice DICL.
0.5
0.4
5
10
15
20
n_com ponents
C.2 INDEPENDENT COMPONENT ANALYSIS (ICA)
DICL-(s.a)
ICA is a statistical and computational technique used to separate a multivariate signal into additive, statistically independent components. Unlike PCA, which decorre- lates the data, ICA aims to find linear transformation that makes the components as independent as possible. Given a data matrix X. ICA assumes that the data is generated linear mixtures of independent components: X AS, where A is un- known mixing matrix and S is the matrix independent components with independent rows. The goal of ICA to estimate unmixing matrix W such that Y WX is good approximation of the independent components S. The implications of ICA on independence are profound: while PCA only guarantees uncorrelated components, ICA goes a step further by optimizing for statistical independence, often measured by con-Gaussyy (kurtosis or negentropy).
Figure 9: Ablation study on the number of principal components in the DICL- (s,a) method.
Fig. 10 shows the estimated mixing matrix A when run- ning ICA on the D4RL-expert dataset on the Hopper en- vironment. Under the assumptions of ICA, notably the statistical independence of the source signals, their lin- ear mixing and the invertibility of the original (unknown) mixing matrix, the original sources are successfully re- covered if each line of the estimated mixing matrix is mostly dominated by a single value, meaning that it's close to an identity matrix up to a permutation with seal- ing. In the case of our states and actions data, it's not clear that this is the case from Fig. 10. Similarly to PCA, can transform the in-context multi-dimensional signal using ICA, and apply the ICL procedure to the recovered independent sources. We plan exploring this method in future follow-up work.
ICA Estimated Mixing matrix
Figure 10: ICA estimated mixing ma- trix.
C.3 AUTOENCODER-BAN BASED APPROACH
Variational Autoencoders (VABs) (Kingma & Welling, 2022) offer a powerful framework for learn- ing representations. A disentangled representation one where each dimension of the latent space captures a distinct and interpretable factor of variation in the data. By combining an encoder net- work that maps inputs to a probabilistic latent space with decoder network that reconstructs the data, VAEs employ the reparameterization trick to enable backpropagation through the sampling process. The key to disentanglement lies in the KL-divergence term of the VAE loss function, which regularizes the latent distribution to be close to a standard normal distribution. Variants such as 3- VAE (Higgins et al., 2017) further emphasize this regularization by scaling the KL-divergence term, thereby encouraging the model to learn more disentangled representation the potential cost of reconstruction quality. Beyond simple VAEs, there exist previous work in the literature that specif- ically aim at learning factorized posterior distribution in the latent space (Kim & Mnih, 2019).
23

---

Published as a conference paper at ICLR 2025
Although this direction looks promising, it strikes different concerns about the learnability of these models in the low data regime considered in our paper.
C.4 SENSITIVITY ANALYSIS
The preceding analysis examines state dimensions as features within representation space, disre- garding their temporal nature and our ultimate objective of predicting the next state. In practice, our interest lies in capturing the dependencies that most significantly influence the next state through the dynamics function of the MDP. To achieve this, we use Sensitivity Analysis (SA) to investigate how variations in the input of the dynamics function impact its output.
Sensitivity Analysis. Sensitivity analysis is a system- atic approach to evaluate how the uncertainty in the out- put of model can be attributed to different sources of un-
Hopper
certainty in the model's inputs. The One-at-a-Time (OAT) method is a technique used to understand the impact of individual input variables the output of model. In the context of transition function of MDP, the OAT method involves systematically varying one current state or action dimension at time, while keeping all others fixed, and observing the resulting changes the out- put dimensions: 0(sr+1) and 0(sc-1)e where (sx)s. (at); d(a=); and (s++1) ) denote the i-th dimension of the state, the j- th dimension of the action, and the k-th dimension of the next state, respectively.
rootz
rooty
thigh
(eg
66
foot
rootx
rootz
rooty
thigh
2
leg
foot
O(s:);
Figure 11: Sensitivity matrix.
In practice, we measure the sensitivity by applying a perturbation (of scale 10%) to each input dimension separately, reporting the absolute change that occurs each dimension of the output. Precisely, for a deterministic transition function f, input state dimension i, and output dimension k, we measure f( a, f(s, a)&l where e, 0.1 scale(t) and elsewhere. The sensitivity matrix in Fig. 11 demonstrates that most of the next state dimensions are mostly affected by their respective previous values (the diagonal shape in the state dimensions square). addition that, actions only directly affect some state dimensions, specifically velocities, which is expected from the nature the physics simulation underlying those systems. This finding suggests that the vICL method might give good results in practice for the considered RL environments, and makes us hope that the DICL-(s) approach is enough to capture the state dimensions dependencies, especially for single-step prediction.
Remark C.I. This sensitivity analysis is specific to the single-step transition function. In practice, such conclusions might change when looking at larger time scale of the simulation.
D ALGORITHMS
D.1 SOFT-ACTOR CRITIC
Soft Actor-Critic (SAC) (Haarnoja et al., 2018) an off off-policy algorithm that incorporates the max- imum entropy framework, which encourages exploration by seeking to maximize the entropy of the policy addition to the expected return. SAC uses deep neural network to approximate the policy (actor) and the value functions (critics), employing two Q-value functions mitigate positive bias in the policy improvement step typical of off-policy algorithms. This approach helps learning more stable and effective policies for complex environments, making SAC particularly suitable for tasks with high-dimensional, continuous action spaces.
We the implementation provided in CleanRL (Huang al., 2022) for SAC. In all environments, we keep the default hyperparameters provided with the library, except for the update frequency. We specify in Table the complete list of hyperparameters used for every considered environment.
24

---

Published as a conference paper at ICLR 2025
Table 2: SAC hyperparameters.
Environment
HalfChectah
Hopper
Pendulum
Update frequency
1000
1000
200
Learning starts
5000
1000
Batch size
128
128
61
Total timesleps
le6
1c6
Gamma
0.99
policy learning rale
Se
4
3e 4
Se 4
D.2 DICL-SAC
For algorithm, integrate LLM inference interface (typically the Transformers library from Huggingface (Wolf et al., 2020)) with CleanRL (Huang et al., 2022). Table 3 shows all DICL-SAC hyperparameter choices for the considered environments.
Table 3: DICL-SAC hyperparameters.
Environment
HalfChectah
Hopper
Pendulum
Update frequency
1000
1000
200
Learning starts
5000
5000
1000
LIM Learning
10000
10000
2000
LLM Learning Trequency
256
256
16
Batch
L.M Batch size (o %
(55%), (3(10%)..2(255)))
7(5%),
(3(10%)) 32(25%)
4(5%).. 7(10%). 16(25%)
timestops
le4
Gamma
0.99
0.99
0.99
Max context length
500
500
198
mode Min context length
LLM sampling method
mode
mode
dynamics learner
Balancing gradient updates. To ensure that DICL-SAC performs equally important gradient up- dates on the LLM generated data, we used a gradient updates balancing mechanism. Indeed, since the default reduction method of loss functions is averaging, the batch B with the smallest batch size gets assigned a higher weight when doing gradient descent: "BT
To address this, we multiply
the loss corresponding the LLM generated batch Bum with correcting coefficient equal weighting across all samples.
[ad ensuring B
We now show the full training curves on the HalfCheetah and Hopper environments (Fig. 12). The return curves show smoothed average training curves = 95% Gaussian confidence intervals for seeds HalfCheetah and Hopper, and 10 seeds for Pendulum.
1.25
le4
HalfCheetah
1e3
Hopper
1.00
3
Return
0.75
Return
2
0.50
0.25
1
0.00
0.0
0.5
1.0
0.0
0.5
1.0
Step
1e6
Step
1e6
SAC
DICL-SACIa 5%)
DICL-SAC(8 10%)
DICL-SAC (a- 25%)
Figure 12: Data-augmented off-policy RL. Full training curves.
25

---

Published as a conference paper at ICLR 2025
The update frequency. The default update frequency of SAC is step, meaning that the policy that interacts with the environment gets updated after every interaction. In our LLM-based framework, this introduces additional layer of complexity this implies that the state visitation distribution of the in-context trajectories will be moving from one timestamp to another. We therefore assume an update frequency equal to the maximal number of steps of episode of given environment.
important to mention that the choice of setting the update frequency for all algorithms to the number of steps equivalent to full episode has dual implications: it stabilize the data collection policy, which is beneficial, but it may also lead to overtraining on data gathered by early, low- quality policies, which is detrimental. This trade-off has been previously studied in the RL literature (Matsushima et al., 2021; Thomas et al., 2024). Notably, Thomas et al. (2024) argues that the update frequency more of a system constraint than design choice or hyperparameter. For instance, controlling a physically grounded system, such as helicopter, inherently imposes a minimal update frequency. Therefore, we deem a fair comparison as this constraint is uniformly applied to all algorithms.
For the sake completeness and comparison, we also evaluated the SAC baseline using its default update frequency of one step. Fig. 13 shows the comparison of our algorithm DICL-SAC, the baseline SAC with update frequency 1000, and the default SAC with update frequency 1. We see that on Halfcheetah the default SAC (uf = 1) performs similarly to SAC with update frequency of 1000. On Pendulum and Hopper it performs slightly better with DICL remaining competitive while having the constraint of an update frequency of 1000.
1e3
Pendulum
1e3 HalfCheetah
le3
Hopper
0.0
2.0
-0.5
Return
Return
Return
1.0
-1.0
0.5
-1.5
0.0
0.0
0.5 Step
1.0
0.25
0.50
0.75
1.00
0.25
0.50
0.75
1.00 1e5
1e4
Step
1e5
Step
SAC(uf 1)
SAC
DICL SAC(a 5%)
DICL
SAC 10%)
DICL - SAC(c 25%)
Figure 13: Data-augmented off-policy RL. Comparison with SAC in the default update frequency regime. We conducted this experiment using the Llama 3.2-IB model.
E WHAT IS THE IMPACT OF THE POLICY ON THE PREDICTION ERROR?
In this experiment, We investigate how policy impacts the accuracy and calibration of our LLM- based dynamics models. To so, we train three model-free algorithms (PPO (Schulman et al., 2017), SAC (Haarnoja al., 2018), and TD3 (Fujimoto et al., 2019)) the HalfCheetah envi- ronment, selecting different checkpoints throughout training to capture diverse policies. We then analyze the correlation between policy characteristics, specifically state coverage (defined as the maximum distance between any two states encountered by the policy) and entropy, with the Mean Squared Error and Kolmogorov-Smimov (KS) statistic. Our findings indicate that the state cover- age correlates with both MSE and KS, possibly because policies that explore a wide range of states generate trajectories that are more difficult to learn. Regarding the entropy, we can see that also correlates with MSE, but interestingly, it does not appear to impact the calibration.
F MULTI-STEP PREDICTION ERRORS
The average multi-step error. In Fig. 4a, compute the average Mean Squared Error over prediction horizons for h = 1,..., 20, and trajectories sampled uniformly from the D4RL expert dataset. For visualization purposes, first rescale all the dimensions (using pipeline composed
26

---

Published as a conference paper at ICLR 2025
1.25
1.25
0.75
0.75
0.5
0.5
0.25
0
1e-1
1.0
1e-1
ppo
0.9
sac
td3
0.8
10
state_coverage
entropy
Correlation plots between state coverage and entropy of policies with MSE and KS metrics under the vICL dynamics learner.
Figure 14:
of MinMaxScaler and StandardScaler) that the respective MSEs the scale. The MSE metric in Table 1 is also computed in a similar fashion, with the exception that it's average over different tasks (HalfCheetah: random, medium, expert; Hopper: medium, expert; Walker2d: medium, expert).
The MLP baseline. For the MLP baseline, we instantiate an MLP with: 4 layers, 128 neurons each, and LU activations. We then format the in-context trajectory as a dataset of ((8t: dr, 8t+1)) on which we train the MI.P for 150 epochs using early stopping and the Adam optimizer (Kingma & Ba, 2015).
We now extend Fig. 4 to show the multi-step generated trajectories for all the dimensions of the HalfCheetah system in Fig. 15,
G CALIBRATION
The naive baseline. In the calibration plots Figs. 8 and 16, we compare the LLM-based dynamics models with a (naive) baseline that estimates a Gaussian distribution using the in-context moments (mean and variance).
KOL.MOOORVV-SMIROO STATISTIC (KS): This metric computed using the quantiles (under the model distribution) of the ground truth values. Hypothetically, these quantiles uniform if the error in predicting the ground truth is a random variable distributed according to a Gaussian with the predicted standard deviation, property the Kolmogorov-S -mirnov (KS) statistics. Formally, starting from the model cumulative distribution function (CDF) (s (18,, 0z), define the empirical CDF of the quantiles of ground truth values
characterize as calibration. To assess this, compute
a.)a for [0, 1]. We denote by Uxx) the CDF of the uniform distribution over the interval [0, 1J. and we define the KS statistics as the largest absolute difference between the two CDFs across the dataset D:
by (z) (ED1
.V
KS(D;0;j (1,...,,d,1)
max ve (1......N) Fo - UUFF( ax))) (10)
The KSS. score ranges between and one, with lower values indicating better calibration.
27

---

Published as a conference paper at ICLR 2025
rootz
rooty
bthigh
bshin
bloot
fthigh
fshin
ffoot
rootx
rootz
rooty
bthigh
bshin
bfoot
fthigh
fshin
ffoot
50
100
150
200
250
300
350
400
groundtruth
VICL
DICLL(()
DICL-(s, a)
MLP
Figure 15: Halfcheetah
H ON THE CHOICE OF THE LLM
In this ablation study, we investigate the impact of LLM size on prediction performance and cali- bration on D4RL. tasks. The LLMs analyzed are all from the I.LaMA 3 family of models (Dubey & the Llama 3 team., 2024), with size range from 1B to 70B parameters, including intermediate sizes of 3B and 8B. Each model is fed with 5 randomly sampled trajectories of length the D4RL datasets: expert, medium, and random. This latter task is only evaluated on HalfCheetah, since the Hopper and Walker2d environments random policies episodes do not have enough context yet apply DICL. For the medium and expert datasets, evaluate them all the environments HalfCheetah (Fig. 17). Hopper (Fig. 18), and Walker2d (Fig. 18b). The metrics used to evaluate the models are:
300 from
28

---

Published as a conference paper at ICLR 2025
rootz
rooty
bthigh
proportion
1.0
0.5
0.0
60
bshin
bfoot
fthigh
proportion
1.0
0.5
0.0
02:13
502
fshin
ffoot
rootx
oroportion
1.0
0.5
...0
rootz
rooty
bthigh
proportion
1.0
0.5
0.0
bshin
bfoot
fthigh
proportion
1.0
0.5
0.0
00330
fshin
ffoot
Sroportion
1.0
0.5
0.0
0.0
0.5
1.0
0.0
0.5 quantile
1.0
quantile
00%
DRUM
Figure 16: Halfcheetah
Mean Squared Error (MSE): Applied after rescaling the data similarly to Appendix F to measure the prediction error.
Dolmogorov-Smimoo (KS) statistic: To evaluate calibration, indicating how well the pre- dicted probabilities match the observed outcomes. This metric is formally described in Appendix G.
All results are averaged over prediction horizons E 20). In the HalfCheetah environment, we observe that DICL-(s) consistently outperforms the other variants across all tasks and with al- most all LLMs in terms of prediction error. DICL-(s, a) is outperformed by vICL in the random and medium datasets, while performance improves the expert dataset. This is likely because the policy has converged to stable expert policy, making it easier for DICL-(s, a) to predict actions as well. Regarding calibration, the three methods generally perform similarly, with a slight advantage for DICL-(s, a), especially with smaller LLMs. In the Hopper environment, the MSE improvement of DICL over vICL is less pronounced with the smallest LLMs but becomes more evident with the LLaMA 3.1 70B model. However, DICL-(s. a) consistently and significantly outperforms both vICL and DICL-(s) terms the KS statistic (calibration). In the Walker2d environment, vICL proves to be strong baseline in the expert task, while DICL-(s) shows improvements over it in
29

---

Published as a conference paper at ICLR 2025
the medium dataset. For calibration in Walker2d, DICL- a) continues to outperform the other variants across all tasks and LLM sizes.
HalfCheetah expert
HalfCheetah medium
HalfCheetah random
MSE
55S
MSE
method
DICL-(s,a) a)
DICL-(s)
vICL
3.2-18
3.2-38
3.1-88
3-88
3.1 3.1-708
3.2-18
3.2-38
3.1-88
3-88
21-708
3.2-18
3.2-8
3.1.08
3.80
3.1-708
llm
llm
llm
HalfCheetah random
HaifCheetah expert
HalfCheetah medium
0.4
0.2
0.3
2
0.2
12
method
0.1
DICL-(s,
0.1
DICL-(s)
vICL
0.0
3,2-18
3.2-38
3.1-88
3.89
3.1-708
3
3.2-88
3.1-80
3.80
3.1-708
3
3.2-38
3.1-00
3-80
3.1-708
llm
llm
llm
Figure 17: HalfCheetah.
Hopper- expert
Hopper- medium
Walker2d expert
Walker2d medium
method DICL-Is, a)
method
DICL
vILL
DICL-tr)
DICL-ss
2.2.10
3.2.0
3.148
3:0
311770
22:18
3.239
3.440
3-a0
3.17708
3.248
3.2-0
2.45
3408
3119108
22:11
3.238
3.1400
2.00
3.17709
Ilm
llm
0.8
Hopper- expert
0.8
Hopper medium
Walker2d export
Wolker2d medium
0.6
0.6
0.6
0.6
method
0.4
method
0.4
vKL
DICL-(s)
0.2
0.2
0.
VIEL
DICL-(s)
DICL-1s, a)
32:16
92:38
3.148
1.56
0.0
2.241
3.239
3.2.8
32.38
2.50
31171
0.0
3.240
22330
lm
lim
(a) Hopper.
(b) Walker2d.
30


Reviews:
[{"review_id": "j06EBQCa5I", "reviewer": "Reviewer_Um7G", "summary": "This paper proposes a new approach Disentangled In-Context Learning (DICL) to generalize LLM-based in-context learning to the domain of continuous-state-space reinforcement learning. This paper then analyzed the theoretical properties of this new DICL approach, and reports empirical experiments results from policy evaluation and data-augmented off-policy RL to the advantages in sample efficiency and performance of DICL in comparison to baseline methods.", "questions": "1. Typo - on Line 63, should it be ‘deferred’ instead of ‘differed’?\n2. Typo - on Line 322, should it be ‘The goal is to improve …’?\n3. In the DICL-SAC algorithm, what would be the optimal value for \\alpha? Are there any intuitions for choosing the optimal \\alpha value?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper has good originality in that it innovatively proposes a novel DICL method to generalize ICL to continuous-state-space RL. This paper also includes solid mathematical derivations and proofs and detailed experiment results to support the claims in the paper.", "weaknesses": "There is a lot of room of improvement for the clarity, writing and presentation of this paper. Multiple places in the paper are not very clearly explained and the general theme of the paper is a little bit hard to follow in its writing. For example, throughout the whole paper, there is no explicit explanation or demonstrations on how exactly the LLM prompts for DICL are constructed. One or more concrete prompt examples would be very helpful in the paper to help readers understand the core technical details of the proposed DICL method.\n\nThis paper has very good potential, but its current form could benefit a lot from a systematic revision that improves its clarity and presentation.", "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": "xpx7Utgi0f", "reviewer": "Reviewer_1htQ", "summary": "The paper explores the utilization of LLMs in the realm of RL, specifically targeting continuous state spaces that have been understudied in the context of LLM integration. It introduces a novel approach termed Disentangled In-Context Learning (DICL), designed to leverage pre-trained LLMs for predicting the dynamics of continuous Markov decision processes. The paper is supported by theoretical analysis and extensive experiments, demonstrating the potential of LLMs in model-based policy evaluation and data-augmented off-policy RL, and shows that LLMs can produce well-calibrated uncertainty estimates.", "questions": "Why does Principal Component Analysis (PCA) decouple features, isn’t it primarily for feature selection? If it is for feature selection, are there any special constraints needed for article-related scenarios?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper presents a method to integrate state dimension interdependence and action information into in-context trajectories within RL environments, enhancing the applicability of LLMs in continuous state spaces.\n2. It provides a theoretical analysis of the policy evaluation algorithm resulting from multi-branch rollouts with LLM-based dynamics models, leading to a novel return bound that enhances understanding in this area.\n3. The paper offers empirical evidence supporting the benefits of LLM modeling in two RL applications: policy evaluation and data-augmented offline RL, showcasing the practicality of the proposed methods.\n4. It demonstrates that LLMs can act as reliable uncertainty estimators, a desirable trait for MBRL algorithms.", "weaknesses": "1. The paper does not extensively discuss how the proposed method generalizes across different environments or tasks, that is, more discussion about the application of this method is needed.\n2. While DICL simplifies certain aspects of RL, the integration of actions and the handling of multivariate data present ongoing challenges. More discussion about the introduced aspects of the DICL is needed.\n3. The experiments are somewhat simplistic, and it would be worthwhile to conduct more in-depth analyses, such as discussing when they become ineffective and more results will be beneficial.\n4. The writing style of the paper is a bit convoluted, making it less fluent to read.", "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": "6gVTJvcivR", "reviewer": "Reviewer_D8FE", "summary": "This paper aims to leverage large language models (LLMs) to achieve model-based reinforcement learning (MBRL). Specifically, it explores using the in-context learning (ICL) capability of LLMs to autoregressively predict the next state and reward, denoted as $<s_t, r_t>$, based on a given prior states and actions.\n\nThe authors' first approach, vanilla ICL (vICL), simplifies the objective by predicting $s_t$ solely based on prior states $s_{1:t-1}$, without considering the action. The second approach, DICL, uses Principal Component Analysis (PCA) to map the state-action vector to a latent space, where ICL is applied to the latent vectors. Additionally, they apply their LLM-based dynamics learner to augment the replay buffer of the Soft Actor-Critic (SAC) algorithm.\n\nAnother important contribution of this work is demonstrating the bound on the difference between true dynamics and LLM-based dynamics under a multi-branch rollout setting.", "questions": "n/a", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "(please also see the summary)\n1. This work systematically studies LLM-based MBRL.  \n2. The method is clearly presented and easy to follow.", "weaknesses": "The “zero-shot” claim is potentially misleading. If \"zero-shot\" is defined at the trajectory level, it is true that no trajectory-level examples were shown to the LLM during prediction. However, as shown in Section 4 (theoretical analysis), it appears necessary to use true dynamics to predict the transition and reward for steps $t < T$. These transitions, such as from $<s_{t-1}, a_{t-1}>$ to $s_t$, effectively serve as state-level few-shot examples. In my understanding, a true “zero-shot” setting would require that all previous transitions be predicted by the LLM itself autoregressively.\n\nAnother related concern is that the experimental setup is difficult to understand. For instance, Figure 3b was unclear to me. Given that the authors claim their method is zero-shot, I was unsure why, in the zero-shot setting, the agreement between the ground truth and the LLM-based method improves over time significantly. Wouldn’t error accumulation or distributional shift cause the LLM’s performance to degrade as more time steps are taken? After reviewing the theoretical analysis section, it seems likely that the result in Figure 3b is due to true dynamics being incorporated into the predictions, which was not mentioned. Please let me know if my understanding is inaccurate.\n\nFinally, the ICL of LLM leverages the whole information in the previous context to predict the next transition. However, the MLP method can only accept the information of previous step as the input. The comparison is somehow unfair.", "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": "M94cCDNhWW", "reviewer": "Reviewer_fXTe", "summary": "The paper investigates zero-shot reinforcement learning (RL) for continuous control tasks using large language models (LLMs) and introduces \"Disentangled In-Context Learning\" (DICL) to improve LLMs' handling of continuous MDPs. This approach separates state and action features to better utilize in-context learning, with experiments validating its efficacy.", "questions": "You mentioned that you change the update step of SAC to accommodate LLM's requirements. Is this unfair? Have you compared with a baseline SAC with update step 1?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors proposed an innovative integration of LLMs into continuous reinforcement learning tasks.\n2. Clear theoretical foundation supporting the DICL framework is shown.\n3. The authors conduct extensive experiments to verify the effectiveness of the method across diverse RL scenarios.\n4. The authors also did insightful analysis of zero-shot capabilities in complex MDPs.", "weaknesses": "Unclear experimental settings. The authors mentioned that they used 3 Llama3 models for the experiments, but all the experiment figures only show one result without stating which model is used. Also, an ablation study about the effect of different LLMs is missing.", "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)