Zero-shot Model-based Reinforcement Learning using Large Language Models

## Abstract

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 .

## 1 Introduction

The rise of large language models (LLMs) has significantly impacted the field of Natural Language Processing (NLP). LLMs (Brown et al., ; Hugo Touvron & the Llama 2 team., ; Dubey & the Llama 3 team, ), which are based on the transformer architecture (Vaswani et al., ), have redefined tasks such as machine translation (Brown et al., ), sentiment analysis (Zhang et al., ), and question answering (Roberts et al., ; Pourkamali & Sharifi, ) by enabling machines to understand and generate human-like text with remarkable fluency. One of the most intriguing aspects of LLMs is their emerging capabilities, particularly in-context learning (ICL) (von Oswald et al., ).
Through ICL, an LLM can learn to perform a new task simply by being provided examples of the task within its input context, without any gradient-based optimization. This phenomenon has been observed not only in text generation but also in tasks such as image classification (Abdelhamed et al., ; Zheng et al., ) and even solving logic puzzles (Giadikiaroglou et al., ), which is unexpected in the context of the standard statistical learning theory. To our knowledge, ICL capabilities of pre-trained LLMs have been only scarcely explored in reinforcement learning (Wang et al., ) despite the demonstrated success of the former in understanding the behavior of deterministic and chaotic dynamical systems (Liu et al., ).

In this paper, we show how ICL with pre-trained LLMs can improve the sample efficiency of Reinforcement 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 et al. () and experimentally studied in Liu et al. (), we use the observed trajectories of a given agent to predict its future state and reward in commonly used realistic RL 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 LLM’s context and 2) handling the interdependence between the state-actions dimensions, as prior approaches were known to treat multivariate data’s covariates independently.

Our approach leads to several novel insights and contributions, which we summarize as follows:

- 1.
*Methodological.* We develop a novel approach to integrate state dimension interdependence and action information into in-context trajectories. This approach, termed Disentangled In-Context Learning (DICL), leads to a 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 can benefit two RL applications: policy evaluation and data-augmented offline RL. Furthermore, we show that the LLM is a calibrated uncertainty estimator, a desirable property for MBRL algorithms.

#### Organization of the paper.

The paper is structured as follows:
[Section 2](#S2) introduces the main concepts from the literature used in our work (while a more detailed related work is differed to [Appendix B](#A2)).
We then start our analysis in [Section 3.1](#S3.SS1), by analyzing LLM’s attention matrices.
DICL is presented in [Section 3.3](#S3.SS3), while
[Section 4](#S4) contains different applications of the proposed method in RL, along with the corresponding theoretical analysis.

## 2 Background knowledge

Reinforcement Learning (RL). The standard framework of RL is the infinite-horizon Markov decision process (MDP) $\mathcal{M}=\langle\mathcal{S},\mathcal{A},P,r,\mu_{0},\gamma\rangle$ where $\mathcal{S}$ represents the state space, $\mathcal{A}$ the action space, $P:\mathcal{S}\times\mathcal{A}\to\mathcal{S}$ the (possibly stochastic) transition dynamics, $r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ the reward function, $\mu_{0}$ the initial state distribution, and $\gamma\in[0,1]$ the discount factor.
The goal of RL is to find, for each state $s\in\mathcal{S}$, a distribution $\pi(s)$ over the action space $\mathcal{A}$, called the policy, that maximizes the expected sum of discounted rewards $\eta(\pi):=\mathbb{E}_{s_{0}\sim\mu_{0},a_{t}\sim\pi,\,s_{t>0}\sim P^{t}}[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})]$. Under a policy $\pi$, we define the state value function at $s\in\mathcal{S}$ as the expected sum of discounted rewards, starting from the state $s$, and following the policy $\pi$ afterwards until termination: $V^{\pi}(s)=\mathbb{E}_{a_{t}\sim\pi,s_{t>0}\sim P^{t}}\big{[}\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})\mid s_{0}=s\big{]}$.

Model-based RL (MBRL). MBRL algorithms address the supervised learning problem of estimating the dynamics of the environment $\hat{P}$ (and sometimes also the reward function $\hat{r}$) from data collected when interacting with the real system.
The model’s loss function is typically the log-likelihood $\mathcal{L}(\mathcal{D};\hat{P})=\frac{1}{N}\sum_{i=1}^{N}\log\hat{P}(s^{i}_{t+1}|s^{i}_{t},a^{i}_{t})$ or Mean Squared Error (MSE) for deterministic models.
The learned model can subsequently be used for policy search under the MDP $\widehat{\mathcal{M}}=\langle\mathcal{S},\mathcal{A},\hat{P},r,\mu_{0},\gamma\rangle$.
This MDP shares the state and action spaces $\mathcal{S},\mathcal{A}$, reward function $r$, with the true environment $\mathcal{M}$, but learns the transition probability $\hat{P}$ from the dataset $\mathcal{D}$. The policy $\hat{\pi}=\operatorname*{arg\,max}_{\pi}\hat{\eta}(\pi)$ learned on $\widehat{\mathcal{M}}$ is not guaranteed to be optimal under the true MDP $\mathcal{M}$ due to distribution shift and model bias.

Large Language Models (LLMs). Within the field of Natural Language Processing, Large Language Models (LLMs) have emerged as a powerful tool for understanding and generating human-like text.
An LLM is typically defined as a neural network model, often based on the transformer architecture (Vaswani et al., ), that is trained on a vast corpus of sequences, $U=\{U_{1},U_{2},\ldots,U_{i},\ldots,U_{N}\}$, where each sequence $U_{i}=(u_{1},u_{2},\ldots,u_{j},\ldots,u_{n_{i}})$ consists of tokens $u_{j}$ from a vocabulary $\mathcal{V}$.
Decoder-only LLMs (Radford et al., ; Dubey & the Llama 3 team, ) typically encode an autoregressive distribution, where the probability of each token is conditioned only on the previous tokens in the sequence, expressed as $p_{\theta}(U_{i})=\prod_{j=1}^{n_{i}}p_{\theta}(u_{j}|u_{0:j-1})$.
The parameters $\theta$ are learned by maximizing the probability of the entire dataset, $p_{\theta}(U)=\prod_{i=1}^{N}p_{\theta}(U_{i})$. Every LLM has an associated tokenizer, which breaks an input string into a sequence of tokens, each belonging to $\mathcal{V}$.

Figure: Algorithm 1 $\text{ICL}_{\theta}$ (Liu et al., ; Gruver et al., )

In-Context Learning (ICL). In order to use trajectories as inputs in ICL, we use the tokenization of time series proposed in Gruver et al. () and Jin et al. (). The latter works use a subset of the LLM sub-vocabulary $\mathcal{V}_{\text{num}}$ representing digits as summarized in [Algorithm 1](#alg1). Given an univarite time series, we first rescale it into a specific range (Liu et al., ; Zekri et al., ; Requeima et al., ) then encode the rescaled time series with $k$ digits. The logits corresponding to tokens in $\mathcal{V}_{\text{num}}$ can be further used to predict the next timestamp transition rule for Markovian systems (Liu et al., ).

## 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 a pre-trained LLM model is sensitive to the primitive univariate signals akin to those encountered in them. For this, we investigate the attention mechanism of the Llama3 8B model (Dubey & the Llama 3 team, ) when we feed it with different signals, including the periodic *fthigh* dimension from the HalfCheetah system (Brockman et al., ).
By averaging the attention matrices over the 32 heads for each of the 32 layers of the multi-head attention in Llama3, we observed distinct patterns that provide insight into the model’s focus and behavior ([Fig. 1](#S3.F1) shows selected attention layers for each signal).
The attention matrices exhibit a diagonal pattern, indicative of strong self-correlation among timestamps, and a subtriangular structure due to the causal masked attention in decoder-only transformers.

Figure: Figure 1: LLM can perceive time patterns. The LLM is fed with $3$ time series presenting distinct patterns. (a) Rectangular pulse. (b) Rectangular signal with constant sub-parts. (c) The *fthigh* dimension of HalfCheetah under an expert policy. Tokens belonging to constant slots (or peaks) attend to all the similar ones that precede them, yet they focus more on the first occurrence of that phenomena.
Refer to caption: /html/2410.11711/assets/x1.png

Further examination of the attention matrices reveals a more intricate finding.
Tokens within repeating patterns (e.g., signal peaks, constant parts) not only attend to past tokens within the same cycle but also to those from previous occurrences of the same pattern, demonstrating a form of in-context learning. The ability to detect and exploit repeating patterns within such signals is especially valuable in RL, where state transitions and action outcomes often exhibit cyclical or recurring dynamics, particularly in continuous control tasks. However, applying this insight to RL presents two critical challenges related to 1) the integration of actions into the forecasting process, and 2) handling of 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 an initial trajectory $\mathcal{T}=(s_{0},a_{0},r_{1},s_{1},a_{1},r_{2},s_{2},\ldots,r_{T-1},s_{T-1})$ of length $T$, with $s_{t}\in\mathcal{S}$, $a_{t}=\pi(s_{t})\in\mathcal{A}$(^1^11In practice, states and actions are real valued vectors spanning a space of dimensions respectively $d_{s}$ and $d_{a}$: $\mathcal{S}=\mathbb{R}^{d_{s}}$, $\mathcal{A}=\mathbb{R}^{d_{a}}$), where the policy $\pi$ is fixed for the whole trajectory, and $r_{t}\in\mathbb{R}$, we want to predict future transitions: given $(s_{T-1},a_{T-1})$ predict the next state and reward $(s_{T},r_{T})$ and subsequent transitions autoregressively. For simplicity we first omit the actions and the reward, focusing instead on the multivariate sequence $\tau^{\pi}=(s_{0},s_{1},\ldots,s_{T})$ where we assume that the state dimensions are independent. Later, we show how to relax the assumptions of omitting actions and rewards, as well as state independence, which is crucial for applications in RL. The joint probability density function of $\tau^{\pi}$ can be written as:

$$ $\begin{cases}\mathbf{P}(\tau^{\pi})=\mu_{0}(s_{0})\prod_{t=1}^{T}P^{\pi}(s_{t}|s_{t-1})\[5.0pt] \text{where }P^{\pi}(s_{t}|s_{t-1})=\int_{a\in\mathcal{A}}\pi(a|s_{t-1})P(s_{t}|s_{t-1},a)\,da\enspace.\end{cases}$ (1) $$

Using the decoder-only nature of the in-context learner defined in [Section 2](#S2), we can apply [Algorithm 1](#alg1) to each dimension of the state vector to infer the transition rule of each visited state in $\tau^{\pi}$ conditioned on its relative history: for all $j\in\{1,\ldots,d_{s}\}$,

$$ $\{\hat{P}^{\pi,j}_{\theta}(s^{j}_{t}|s^{j}_{t-1},\ldots,s^{j}_{1},s^{j}_{0})\}_{t\leq T}=\text{ICL}_{\theta}(\tau^{\pi,j})$ (2) $$

where $\theta$ are the fixed parameters of the LLM used as an in-context learner, and $T$ its context length. Assuming complete observability of the MDP state, the Markovian property unveils an equivalence between the learned transition rules and the corresponding Markovian ones: $\hat{P}_{\theta}(s_{t}|s_{t-1},\ldots,s_{1},s_{0})=\hat{P}_{\theta}(s_{t}|s_{t-1})$.

This approach, that we name vICL (for vanilla ICL), thus applies [Algorithm 1](#alg1) on each dimension of the state individually, assuming their independence. Furthermore, the action information is integrated-out (as depicted in [Eq. 1](#S3.E1)), which in theory, limits the application scope of this method to quantities that only depend on a policy through the expectation over actions (e.g., the value function $V^{\pi}(s)$). We address these limitations in the next section.

### 3.3 State and action dimension interdependence

Figure: Figure 2: The covariance matrix computed from the D4RL expert dataset in the Halfcheetah environment indicates linear correlations between state-action features. We perform PCA before applying ICL to achieve a potentially lower-dimensional space with linearly uncorrelated features.
Refer to caption: /html/2410.11711/assets/x2.png

In this section we address the two limitations of vICL discussed in [Section 3.2](#S3.SS2) by introducing Disentangled In-Context Learning (DICL), a method that relaxes the assumption of state feature independence and reintroduces the action by employing strategies that aim to map the state-action vector to a 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 a speed-up of the overall approach.

While sophisticated approaches(^2^22A more detailed discussion of alternative approaches to PCA is provided in [Appendix C](#A3).) like disentangled autoencoders could be considered for DICL, in this work we employ Principal Component Analysis (PCA).
In fact, the absence of pre-trained models for this type of representation learning requires training from scratch on a potentially large dataset.
This goes against our goal of leveraging the pre-trained knowledge of LLMs and ICL.
Instead, we find that PCA, which generates new linearly uncorrelated features and can reduce dimensionality, strikes a good balance between simplicity, tractability, and performance ([Fig. 2](#S3.F2) and [Fig. 3](#S3.F3)).
Nonetheless, DICL is agnostic to this aspect and any approach that can disentangle features can be used in place of PCA.

In the rest of the paper, we refer to the routine that applies PCA to feature space of states and actions and then runs [Algorithm 1](#alg1) in the projection space of principal components as DICL-$(s,a)$. In some cases, integrating the action may not be necessary, such as when one is only interested in estimating the value function $V^{\pi}(s)$. To address this, we introduce DICL-$(s)$ which denotes the combination of PCA and ICL applied solely to the states.

### 3.4 An illustrative example

Figure: ((a)) Multi-step error.
Refer to caption: /html/2410.11711/assets/x3.png

In this section, we aim to challenge our approach against the HalfCheetah system from the MuJoCo Gym environment suite (Brockman et al., ; Todorov et al., ).
All our experiments are conducted using the Llama 3 series of models (Llama 3 8B, Llama 3.2 1B, Llama 3.2 3B) (Dubey & the Llama 3 team, ).
[Fig. 3(a)](#S3.F3.sf1) shows the average MSE over a prediction horizon of $h\in\{1,\ldots,20\}$ steps for each state dimension. [Fig. 3(b)](#S3.F3.sf2) shows predicted trajectories for selected state dimensions of the HalfCheetah system (the details of the experiment, the metrics and the remaining state dimensions are differed to [Appendix F](#A6)).

We first observe that the LLM-based dynamics forecasters exhibit a burn-in phase ($\approx 70$ steps in [Fig. 3(b)](#S3.F3.sf2)) that is necessary for the LLM to gather enough context. For multi-step prediction, [Fig. 3(a)](#S3.F3.sf1), 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 dimensionality reduction feature by selecting a number of components equal to half the number of the original features.
This results in a significant decrease in the computational time of the method without loss of performance, as showcased by [Fig. 3(c)](#S3.F3.sf3). We now move to integrating our methodology in RL-related applications.

## 4 Use-cases in Reinforcement Learning

As explored in the preceding sections, LLMs can be used as accurate dynamics learners for proprioceptive control through in-context learning.
We now state our main contributions in terms of the integration of DICL into MBRL.
First, we generalize the return bound of Model-Based Policy Optimization (MBPO) (Janner et al., ) to the more general case of multiple branches and use it to analyze our method.
Next, we leverage the LLM to augment the replay buffer of an off-policy RL algorithm, leading to a more sample-efficient algorithm.
In a second application, we apply our method to predict the reward signal, resulting in a hybrid model-based policy evaluation technique.
Finally, we show that the LLM provides calibrated uncertainty estimates and conclude with a discussion of our results.

### 4.1 Theoretical analysis: Return bound under multi-branch 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 $\gamma$), the prediction horizon $k$, and the model generalization error $\varepsilon_{\text{m}}$. As established in Janner et al. () and Frauenknecht et al. (), the framework for deriving these theoretical guarantees is the one of branched model-based rollouts.

Figure: Figure 4: Multi-branch return. The rollout following the true dynamics $P$ is shown in blue. The branched rollouts following LLM-based dynamics $\hat{P}_{\text{llm}}$ are in purple. Branched rollouts can overlap and the corresponding return is the expectation over the overlapping branches.
Refer to caption: /html/2410.11711/assets/x7.png

A branched rollout return $\eta^{\text{branch}}[\pi]$ of a policy $\pi$ is defined in Janner et al. () 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 $\hat{P}$ for $k$ steps.
For our LLM-based dynamics learner, we are 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.

###### Definition 4.1 (Multi-branch rollout return) .

The multi-branch rollout return $\eta^{\text{llm}}_{p,k,T}[\pi]$ of a policy $\pi$ is defined as the expected return over rollouts with the following dynamics:

- 1.
for $t<T$, where $T$ is the minimal context length, the rollout follows the true dynamics $P$.
- 2.
for $t\geq T$, with probability $p$, the rollout switches to the LLM-based dynamics $\hat{P}_{\text{llm}}$ 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. [4](#S4.F4)).

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 trajectories, $p\in[0,1]$ the probability that a given state is a 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 a policy $\pi$ is bounded at each timestep by $\max_{t\geq T}\mathbb{E}_{s\sim P^{t},a\sim\pi}[D_{\text{TV}}(P(.|s,a)||\hat{P}_{\text{llm}}(.|s,a))]\leq\varepsilon_{\text{llm}}(T)$. Then under a multi-branched rollout scheme with a branch length of $k$, the return is bounded as follows:

$$ $|\eta(\pi)-\eta^{\text{llm}}_{p,k,T}(\pi)|\leq 2\frac{\gamma^{T}}{1-\gamma}r_{\max}k^{2}\,p\,\varepsilon_{\text{llm}}(T)\enspace,$ (3) $$

where $r_{\max}=\max_{s\in\mathcal{S},a\in\mathcal{A}}r(s,a)$.

[Theorem 4.2](#S4.Thmtheorem2) generalizes the single-branch return presented in Janner et al. (), incorporating an additional factor of the prediction horizon $k$ due to the presence of multiple branches, and directly accounting for the impact of the amount of LLM training data through the branching factor $p$. Additionally, the bound is inversely proportional to the minimal context length $T$, both through the power in the discount factor $\gamma^{T}$ and the error term $\varepsilon_{\text{llm}}(T)$. Indeed, the term $\varepsilon_{\text{llm}}(T)$ corresponds to the generalization error of in-context learning. Several works in the literature studied it and showed that it typically decreases in $\mathcal{O}(T^{-1/2})$ with $T$ the length of the context trajectories (Zekri et al., ; Zhang et al., ; Li et al., ).

###### Definition 4.1 (Multi-branch rollout return) .

###### Theorem 4.2 (Multi-branch return bound) .

### 4.2 Data-augmented off-policy Reinforcement Learning

In this section, we show how DICL can be used for data augmentation in off-policy model-free RL algorithms such as Soft Actor-Critic (SAC) (Haarnoja et al., ).
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 to improve sample-efficiency and accelerate the learning curve, particularly in the early stages of learning as the LLM can generate accurate transitions from a small trajectory.
We name this application of our approach DICL-SAC.

Figure: Algorithm 2 DICL-SAC

As defined in Corrado & Hanna (), data-augmented off-policy RL involves perturbing previously observed transitions to generate new transitions, without further interaction with the environment. The generated transitions should ideally be diverse and feasible under the MDP dynamics to enhance sample efficiency while ensuring that the optimal policy remains learnable. [Algorithm 2](#alg2) (DICL-SAC) integrates multiple components to demonstrate a novel proof-of-concept for improving the sample efficiency of SAC using DICL for data augmentation. Let
$\mathcal{T}=(s_{0},a_{0},r_{0},\ldots,s_{T_{\max}},a_{T_{\max}},r_{T_{\max}})$ be a real trajectory collected with a fixed policy $\pi_{\phi}$, sampled from the real transitions being stored in a replay buffer $\mathcal{R}$. We generate synthetic transitions $(s_{t},\tilde{a}_{t},r_{t},\hat{s}_{t+1})_{T\leq t\leq T_{\max}}$, where $\hat{s}_{t+1}$ is the next state generated by the LLM model applied on the trajectory of the states only, $\tilde{a}_{t}$ is an action sampled from the data collection policy $\pi_{\phi}(.|s_{t})$, and $T$ is the minimal context length.
These transitions are then stored in a separate replay buffer $\mathcal{R}_{\text{llm}}$.
At a given update frequency, DICL-SAC performs $G$ gradient updates using data sampled from $\mathcal{R}$ and $\alpha\%\cdot G$ gradient updates using data sampled from $\mathcal{R}_{\text{llm}}$.
Other hyperparameters of our method include the LLM-based method (vICL or DICL-$(s)$), how often we generate new LLM data and the maximal context length $T_{\max}$ (see [Appendix D](#A4) for the full list of hyperparameters).

It is possible to use the LLM as a transition model with DICL-$(s,a)$, generating the next state based on the previous state and an action (potentially sampled from a different policy), this approach would make data generation very expensive due to the high inference cost of the LLM. Indeed, thanks to [Eq. 2](#S3.E2), our data augmentation strategy, applies the LLM only to the states, allowing us to generate $T_{\max}-T$ transitions with $d_{s}$ calls to the LLM, one call for each dimension of the state. However, using the LLM as a transition model, accounting for actions, would require $(T_{\max}-T)\cdot d_{s}$ calls as we would need one call for each action we modify. This is also why we assume a fixed policy in the context, enabling the LLM to learn $P^{\pi_{\phi}}$ using only the states.

Figure: Figure 5: Data-augmented off-policy RL. In the early stages of training DICL-SAC improves the sample efficiency of SAC on three Gym control environments.
Refer to caption: /html/2410.11711/assets/x8.png

[Fig. 5](#S4.F5) compares the return curves obtained by DICL-SAC against SAC in three control environments from the Gym library (Brockman et al., ). As anticipated with our data augmentation approach, we observe that our algorithm improves the sample efficiency of SAC at the beginning of training. This improvement is moderate but significant in the Pendulum and HalfCheetah environments, while the return curves tend to be noisier in the Hopper environment. Furthermore, as the proportion of LLM data $\alpha$ increases, the performance of the algorithm decreases (particularly in HalfCheetah), as predicted by [Theorem 4.2](#S4.Thmtheorem2). Indeed, a larger proportion of LLM data correlates with a higher probability of branching $p$, as more branching points will be sampled throughout the training. Regarding the other parameters of our bound in [Theorem 4.2](#S4.Thmtheorem2), we set $T=1$, meaning all LLM-generated transitions are added to $\mathcal{R}_{\text{llm}}$, and $k=1$ to minimize LLM inference time.

### 4.3 Policy Evaluation

Figure: Figure 6: Policy evaluation with DICL. Relative error on the predicted value over $k=500$ steps, with context length of $T=500$.
Refer to caption: /html/2410.11711/assets/x12.png

In this section we show how DICL can be used for policy evaluation. System engineers are often presented with several policies to test on their systems. On the one hand, off-policy evaluation (e.g., Uehara et al. ()) involves using historical data collected from a different 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, online evaluation provides a direct and unbiased 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 true performance in production. However, online evaluations can be time-consuming and may temporarily 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. 6](#S4.F6) illustrates the relative error in value obtained by predicting the trajectory of rewards for $k$ steps, given a context length of $T=500$. When $k\leq 500$, we complete the remaining steps of the 1000-step episode using the actual rewards.
For the two versions of DICL, the reward vector is concatenated to the feature space prior to applying PCA.
In the Hopper environment, it is evident that predicting the reward trajectory alone is a challenging task for the vanilla method vICL.
On the contrary, both DICL-$(s)$ and DICL-$(s,a)$ effectively capture some of the dependencies of the reward signal on the states and actions, providing a more robust method for policy evaluation, and matching 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 $\dot{rootx}$ dimension in HalfCheetah, which proved to be harder to predict by our approach, as can be seen in [Fig. 3(a)](#S3.F3.sf1).

Note that the experimental setup that we follow here is closely related to the concept of Model-based Value Expansion (Feinberg et al., ; Buckman et al., ), 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. 3(b)](#S3.F3.sf2) is the confidence interval around the predictions. As detailed in [Algorithm 1](#alg1), one can extract a full probability distribution for the next prediction given the context, enabling uncertainty estimation in the LLM’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 to be an important property of a dynamics model when used in reinforcement learning (Malik et al., ). In this section, we aim to investigate whether the uncertainty estimates derived from the LLM’s logits are well-calibrated. We achieve this by evaluating the quantile calibration of the probability distributions obtained for each LLM-based method.

#### Quantile calibration.

For a regression problem with variable $y\in\mathcal{Y}=\mathbb{R}$, and a model that outputs a cumulative distribution function (CDF) $F_{i}$ over $y_{i}$ (where $i$ indexes data points), quantile calibration implies that $y_{i}$ (groundtruth) should fall within a $p\%$-confidence interval $p\%$ of the time:

$$ $\frac{\sum_{i=1}^{N}\mathbb{I}\{y_{i}\leq F_{i}^{-1}(p)\}}{N}\to p\hskip 8.5359pt\text{for all}\hskip 8.5359ptp\in[0,1]\hskip 8.5359pt\text{as}\hskip 8.5359ptN\to\infty$ (4) $$

where $F_{i}^{-1}:[0,1]\to\mathcal{Y}$ denotes the quantile function $F_{i}^{-1}(p)=\inf\{y:p\leq F_{i}(y)\}$ for all $p\in[0,1]$, and $N$ the number of samples.

Figure: Figure 7: Quantile calibration reliability diagram. We show that the LLM uncertainty estimates are well-calibrated. Vertical lines show the Kolmogorov-Smirnov statistic for each fit.
Refer to caption: /html/2410.11711/assets/x13.png

#### LLMs are well-calibrated forecasters.

[Fig. 7](#S4.F7) 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 to predict the next state, the uncertainty estimates derived from the LLM’s logits are well-calibrated in terms of quantile calibration. Ideally, forecasters should align with the diagonal in [Fig. 7](#S4.F7), which the LLM approach nearly achieves. Furthermore, when comparing with a naive baseline (the details are differed to [Appendix G](#A7)), the LLM-forecaster matches the baseline when it’s already calibrated, and improves over it when it’s not. To quantify a forecaster’s calibration with a point statistic, we compute the Kolmogorov-Smirnov goodness-of-fit test ([Eq. 10](#A7.E10)) that is shown in the legend of [Fig. 7](#S4.F7).

## 5 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 conceptualized in-context learning for univariate time series prediction, and provide a systematic methodology to apply ICL to 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, in 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 our 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 second application, consisted in predicting the trajectory of rewards in order to perform hybrid online and model-based policy evaluation. Finally, we showed that the LLM-based dynamics model also provides well-calibrated uncertainty estimates.

## Reproducibility Statement

In order to ensure reproducibility we release the code at [https://github.com/abenechehab/dicl](https://github.com/abenechehab/dicl). The implementation details and hyperparameters are listed in [Appendix D](#A4).

## Appendix A Theoretical analysis

### A.1 Proof of Theorem 4.2

We start by formally defining the LLM multi-branch return $\eta^{\text{llm}}_{p,k,T}$. To do so, we first denote $A_{t}$ the random event of starting a $k$-step LLM branch at timestep $t$ and we denote $X_{t}$ the associated indicator random variable $X_{t}=\mathbbm{1}[A_{t}]$. We assume that the $(X_{t})_{t\geq T}$ are independent. We then define the random event $A_{t}^{k}$ that at least one of the $k$ preceding timesteps has been branched, meaning that the given timestep $t$ belongs to at least one LLM branch among the $k$ possible branches: $A_{t}^{k}=\bigcup_{i=0}^{k-1}A_{t-i}$. The LLM multi-branch return can then be written as follows:

$$ $\begin{split}\eta^{\text{llm}}_{p,k,T}(\pi)&=\underbrace{\sum_{t=0}^{T-1}\gamma^{t}\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}}_{\text{Burn-in phase to gather minimal context size }T}\\ &+\sum_{t=T}^{\infty}\gamma^{t}\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\bigg{[}\mathbbm{1}[A_{t}^{k}]\underbrace{\frac{1}{\sum_{i=1}^{k}X_{t-i}}\sum_{i=1}^{k}X_{t-i}\mathbb{E}_{s_{t}\sim\hat{P}_{t,\text{llm}}^{i},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}}_{\text{average reward among the branches spanning timestep }t}\\ &+\underbrace{\mathbbm{1}[\bar{A_{t}^{k}}]\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}}_{\text{When no branch is spanning timestep }t}\bigg{]},\end{split}$ (5) $$

where $P^{t}=P(.|P^{t-1})$ with $P^{0}=\mu_{0}$ the initial state distribution and $\hat{P}_{t,\text{llm}}^{i}=\hat{P}_{\text{llm}}^{i}(.|P^{t-i})$.

Before continuing, we first need to establish the following lemma.

###### Lemma A.1 .

(Multi-step Error Bound, Lemma B.2 in Frauenknecht et al. () and Janner et al. ().)
Let $P$ and $\tilde{P}$ be two transition functions. Define the multi-step error at time step $t$, starting from any initial state distribution $\mu_{0}$, as:

$$ $\varepsilon_{t}:=D_{\text{TV}}(P^{t}(\cdot|\mu_{0})\|\tilde{P}^{t}(\cdot|\mu_{0}))$ $$

with $P^{0}=\tilde{P}^{0}=\mu_{0}$.
Let the one-step error at time step $t\geq 1$ be defined as:

$$ $\xi_{t}:=\mathbb{E}_{s\sim P^{t-1}(\cdot|\mu_{0})}\left[D_{\text{TV}}(P(\cdot|s)\|\tilde{P}(\cdot|s))\right],$ $$

and $\xi_{0}=\varepsilon_{0}=0$.

Then, the multi-step error satisfies the following bound:

$$ $\varepsilon_{t}\leq\sum_{i=0}^{t}\xi_{i}.$ $$

###### Proof.

Let $t>0$. We start with the definition of the total variation distance:

$$ $\displaystyle\varepsilon_{t}$ $\displaystyle=D_{\text{TV}}(P^{t}(\cdot|\mu_{0})\|\tilde{P}^{t}(\cdot|\mu_{0}))$ $\displaystyle=\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\left|P^{t}(s^{\prime}|\mu_{0})-\tilde{P}^{t}(s^{\prime}|\mu_{0})\right|\,ds^{\prime}$ $\displaystyle=\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\left|\int_{s\in\mathcal{S}}P(s^{\prime}|s)P^{t-1}(s|\mu_{0})-\tilde{P}(s^{\prime}|s)\tilde{P}^{t-1}(s|\mu_{0})\,ds\right|\,ds^{\prime}$ $\displaystyle\leq\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\int_{s\in\mathcal{S}}\left|P(s^{\prime}|s)P^{t-1}(s|\mu_{0})-\tilde{P}(s^{\prime}|s)\tilde{P}^{t-1}(s|\mu_{0})\right|\,ds\,ds^{\prime}$ $\displaystyle=\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\int_{s\in\mathcal{S}}\left|P(s^{\prime}|s)P^{t-1}(s|\mu_{0})-\tilde{P}(s^{\prime}|s)\tilde{P}^{t-1}(s|\mu_{0})\right|\,ds\,ds^{\prime}$ $\displaystyle=\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\int_{s\in\mathcal{S}}\left|P(s^{\prime}|s)P^{t-1}(s|\mu_{0})-\tilde{P}(s^{\prime}|s)P^{t-1}(s|\mu_{0})\right.$ $\displaystyle\qquad\left.+\tilde{P}(s^{\prime}|s)P^{t-1}(s|\mu_{0})-\tilde{P}(s^{\prime}|s)\tilde{P}^{t-1}(s|\mu_{0})\right|\,ds\,ds^{\prime}$ $\displaystyle\leq\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\int_{s\in\mathcal{S}}P^{t-1}(s|\mu_{0})\left|P(s^{\prime}|s)-\tilde{P}(s^{\prime}|s)\right|\,ds\,ds^{\prime}$ $\displaystyle\qquad+\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\int_{s\in\mathcal{S}}\tilde{P}(s^{\prime}|s)\left|P^{t-1}(s|\mu_{0})-\tilde{P}^{t-1}(s|\mu_{0})\right|\,ds\,ds^{\prime}$ $\displaystyle=\int_{s\in\mathcal{S}}\left[\frac{1}{2}\int_{s^{\prime}\in\mathcal{S}}\left|P(s^{\prime}|s)-\tilde{P}(s^{\prime}|s)\right|\,ds^{\prime}\right]P^{t-1}(s|\mu_{0})\,ds$ $\displaystyle\qquad+\frac{1}{2}\int_{s\in\mathcal{S}}\left(\int_{s^{\prime}\in\mathcal{S}}\tilde{P}(s^{\prime}|s)\,ds^{\prime}\right)\left|P^{t-1}(s|\mu_{0})-\tilde{P}^{t-1}(s|\mu_{0})\right|\,ds$ $\displaystyle=\mathbb{E}_{s\sim P^{t-1}(\cdot|\mu_{0})}\left[D_{\text{TV}}(P(\cdot|\mu_{0})\|\tilde{P}(\cdot|s))\right]+D_{\text{TV}}(P^{t-1}(\cdot|\mu_{0})\|\tilde{P}^{t-1}(\cdot|\mu_{0}))$ $\displaystyle=\xi_{t}+\varepsilon_{t-1}$ $$

Given that $\xi_{0}=\varepsilon_{0}=0$, by induction we have:

$$ $\varepsilon_{t}\leq\sum_{i=0}^{t}\xi_{i}.$ $$

∎

We now restate and prove [Theorem 4.2](#S4.Thmtheorem2):

###### Theorem A.2 (Multi-branch return bound) .

Let $T$ be the minimal length of the in-context trajectories, $p\in[0,1]$ the probability that a given state is a 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 a policy $\pi$ is bounded at each timestep by $\max_{t\geq T}\mathbb{E}_{s\sim P^{t},a\sim\pi}[D_{\text{TV}}(P(.|s,a)||\hat{P}_{\text{llm}}(.|s,a))]\leq\varepsilon_{\text{llm}}(T)$. Then under a multi-branched rollout scheme with a branch length of $k$, the return is bounded as follows:

$$ $|\eta(\pi)-\eta^{\text{llm}}_{p,k,T}(\pi)|\leq 2\frac{\gamma^{T}}{1-\gamma}r_{\max}k^{2}\,p\,\varepsilon_{\text{llm}}(T)\enspace,$ (6) $$

where $r_{\max}=\max_{s\in\mathcal{S},a\in\mathcal{A}}r(s,a)$.

###### Proof.

Step 1: Expressing the bound in terms of horizon-dependent errors.

$$ $\displaystyle|\eta(\pi)-\eta^{\text{llm}}_{p,k,T}(\pi)|$ $\displaystyle=\bigg{\lvert}\sum_{t=T}^{\infty}\gamma^{t}\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}$ $\displaystyle-\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\bigg{[}\mathbbm{1}[A_{t}^{k}]\frac{1}{\sum_{i=1}^{k}X_{t-i}}\sum_{i=1}^{k}X_{t-i}\mathbb{E}_{s_{t}\sim\hat{P}_{t,\text{llm}}^{i},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}$ $\displaystyle-\mathbbm{1}[\bar{A_{t}^{k}}]\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}\bigg{]}\bigg{\lvert}$ $\displaystyle\leq\sum_{t=T}^{\infty}\gamma^{t}\bigg{\lvert}\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\bigg{[}\mathbbm{1}[A_{t}^{k}]\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}+\mathbbm{1}[\bar{A}_{t}^{k}]\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}\bigg{]}$ $\displaystyle-\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\bigg{[}\mathbbm{1}[A_{t}^{k}]\frac{1}{\sum_{i=1}^{k}X_{t-i}}\sum_{i=1}^{k}X_{t-i}\mathbb{E}_{s_{t}\sim\hat{P}_{t,\text{llm}}^{i},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}$ $\displaystyle-\mathbbm{1}[\bar{A_{t}^{k}}]\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}\bigg{]}\bigg{\lvert}$ $\displaystyle\leq\sum_{t=T}^{\infty}\gamma^{t}\bigg{|}\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\bigg{[}\mathbbm{1}[A_{t}^{k}]\bigg{(}$ $\displaystyle\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}-\frac{1}{\sum_{i=1}^{k}X_{t-i}}\sum_{i=1}^{k}X_{t-i}\mathbb{E}_{s_{t}\sim\hat{P}_{t,\text{llm}}^{i},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}\bigg{)}\bigg{]}\bigg{|}$ $\displaystyle\leq\sum_{t=T}^{\infty}\gamma^{t}\bigg{\lvert}\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\bigg{[}\mathbbm{1}[A_{t}^{k}]\frac{1}{\sum_{i=1}^{k}X_{t-i}}\sum_{i=1}^{k}X_{t-i}\bigg{(}$ $\displaystyle\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}-\mathbb{E}_{s_{t}\sim\hat{P}_{t,\text{llm}}^{i},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}\bigg{)}\bigg{]}\bigg{|}$ $$

We then expand the integrals in the terms $\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}-\mathbb{E}_{s_{t}\sim\hat{P}_{t,\text{llm}}^{i},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}$ and express it in terms of horizon-dependent multi-step model errors:

$$ $\begin{split}\mathbb{E}_{s_{t}\sim P^{t},a_{t}\sim\pi}&\big{[}r(s_{t},a_{t})\big{]}-\mathbb{E}_{s_{t}\sim\hat{P}_{t,\text{llm}}^{i},a_{t}\sim\pi}\big{[}r(s_{t},a_{t})\big{]}\\ &=\int_{s\in\mathcal{S}}\int_{a\in\mathcal{A}}r(s,a)\big{(}P^{t}(s,a)-\hat{P}_{t,\text{llm}}^{i}(s,a)\big{)}\;da\;ds\\ &\leq r_{\max}\int_{s\in\mathcal{S}}\int_{a\in\mathcal{A}}\big{(}P^{t}(s,a)-\hat{P}_{t,\text{llm}}^{i}(s,a)\big{)}\;da\;ds\\ &\leq r_{\max}\int_{s\in\mathcal{S}}\int_{a\in\mathcal{A}}\big{(}P^{t}(s)-\hat{P}_{t,\text{llm}}^{i}(s)\big{)}\pi(a|s)\;da\;ds\\ &\leq r_{\max}\int_{s\in\mathcal{S}}\big{(}P^{t}(s)-\hat{P}_{t,\text{llm}}^{i}(s)\big{)}\;ds\\ &\leq 2r_{\max}D_{\text{TV}}(P^{t}||\hat{P}_{t,\text{llm}}^{i})\end{split}$ (7) $$

Step 2: Simplifying the bound.
By applying [Lemma A.1](#A1.Thmtheorem1) we can bound the multi-step errors using the bound on one-step errors:

$$ $D_{\text{TV}}(P^{t}||\hat{P}_{t,\text{llm}}^{i})\leq i\,\varepsilon_{\text{llm}}(T)\leq k\,\varepsilon_{\text{llm}}(T)$ (8) $$

Therefore, the bound becomes:

$$ $\begin{split}|\eta(\pi)-\eta^{\text{llm}}_{p,k,T}(\pi)|&\leq 2r_{\max}\,k\,\varepsilon_{\text{llm}}(T)\sum_{t=T}^{\infty}\gamma^{t}\left|\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\left[\mathbbm{1}[A_{t}^{k}]\frac{1}{\sum_{i=1}^{k}X_{t-i}}\sum_{i=1}^{k}X_{t-i}\right]\right|\\ &=2r_{\max}\,k\,\varepsilon_{\text{llm}}(T)\sum_{t=T}^{\infty}\gamma^{t}\left|\mathbb{E}_{X_{t-i}\sim b(p),1\leq i\leq k}\left[\mathbbm{1}[A_{t}^{k}]\right]\right|\\ &\leq 2r_{\max}\,k\,\varepsilon_{\text{llm}}(T)\sum_{t=T}^{\infty}\gamma^{t}kp\\ &=2\frac{\gamma^{T}}{1-\gamma}r_{\max}k^{2}\,p\,\varepsilon_{\text{llm}}(T)\end{split}$ (9) $$

∎

###### Lemma A.1 .

###### Proof.

###### Theorem A.2 (Multi-branch return bound) .

###### Proof.

## Appendix 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, ; Hafner et al., ; Gal et al., ; Levine & Koltun, ; Chua et al., ; Janner et al., ; Kégl et al., ), and in the offline (pure batch) RL where we do one step of model learning followed by policy learning (Yu et al., ; Kidambi et al., ; Lee et al., ; Argenson & Dulac-Arnold, ; Zhan et al., ; Yu et al., ; Liu et al., ; Benechehab et al., ).
Planning is used either at decision time via model-predictive control (MPC) (Draeger et al., ; Chua et al., ; Hafner et al., ; Pinneri et al., ; Kégl et al., ), or in the background where a model-free agent is learned on imagined model rollouts (Dyna; Janner et al. (); Sutton (); Sutton et al. (); Ha & Schmidhuber ()), or both.
For example, model-based policy optimization (MBPO) (Janner et al., ) 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., ; Yang et al., ), playing key roles in enhancing decision-making (Kannan et al., ; ; Zhang et al., ; Feng et al., ), reward design (Kwon et al., ; Wu et al., ; Carta et al., ; Liu et al., ), and information processing (Poudel et al., ; Lin et al., ). The use of LLMs as world models is particularly relevant to our work. More generally, the Transformer architecture (Vaswani et al., ) has been used in offline RL (Decision Transformer Chen et al. (); Trajectory Transformer Janner et al. ()). Pre-trained LLMs have been used to initialize decision transformers and fine-tune them for offline RL tasks (Shi et al., ; Reid et al., ; Yang & Xu, ). As world models, Dreamer-like architectures based on Transformers have been proposed (Micheli et al., ; Zhang et al., ; Chen et al., ), demonstrating efficiency for long-memory tasks such as Atari games. In text-based environments, LLMs have found multiple applications (Lin et al., ; Feng et al., ; Zhang et al., ; Ma et al., ), including using code-generating LLMs to generate policies in a zero-shot fashion (Liang et al., ; Liu et al., ).

The closest work to ours is Wang et al. (), where a 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. () aim to directly learn a low-level control policy from the LLM, incorporating extra information in the prompt. Furthermore, Wang et al. () 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 1B.

#### 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., ; von Oswald et al., ; Akyürek et al., ; Garg et al., ; Xie et al., ). Regarding time series forecasting, LLMTime (Gruver et al., ) successfully leverages ICL for zero-shot extrapolation of one-dimensional time series data.
Similarly, Das et al. () introduce a foundational model for one-dimensional zero-shot time series forecasting, while Xue & Salim () combine numerical data and text in a question-answer format. ICL can also be used to approximate a continuous density from the LLM logits. For example, Liu et al. () 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. () presented *LLMProcesses*, a method aimed at extracting multi-dimensional distributions from LLMs. Other practical applications of ICL on numerical data include few-shot classification on tabular data (Hegselmann et al., ), regression (Vacareanu et al., ), and meta ICL (Coda-Forno et al., ).

## Appendix C State and action dimensions interdependence - additional materials

### C.1 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 uncorrelated. The principal components can be ordered such that the first few retain most of the variation present in all of the original variables. Formally, given a data matrix $\mathbf{X}$ with $n$ observations and $p$ variables, PCA diagonalizes the covariance matrix $\mathbf{C}=\frac{1}{n-1}\mathbf{X}^{T}\mathbf{X}$ to find the eigenvectors, which represent the directions of the principal components: $\text{PCA: }\mathbf{X}\rightarrow\mathbf{Z}=\mathbf{X}\mathbf{W},\quad\text{where }\mathbf{W}\text{ are the eigenvectors of }\mathbf{C}$. In our case, the data represents a dataset of states and actions given a data collecting policy $\pi_{D}$, while the $p$ variables represent the state (eventually also the action) dimensions.

Figure: Figure 8: Ablation study on the number of principal components in the DICL-$(s,a)$ method.
Refer to caption: /html/2410.11711/assets/x14.png

#### Ablation on the number of components.

[Fig. 8](#A3.F8) shows an ablation study on 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](#A6) 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, before going down again when the full variance is recovered. This finding strengthens the position of PCA as our Disentangling algorithm of choice in DICL.

### C.2 Independent Component Analysis (ICA)

ICA is a statistical and computational technique used to separate a multivariate signal into additive, statistically independent components.
Unlike PCA, which decorrelates the data, ICA aims to find a linear transformation that makes the components as independent as possible.
Given a data matrix $\mathbf{X}$, ICA assumes that the data is generated as linear mixtures of independent components: $\mathbf{X}=\mathbf{A}\mathbf{S}$, where $\mathbf{A}$ is an unknown mixing matrix and $\mathbf{S}$ is the matrix of independent components with independent rows.
The goal of ICA is to estimate an unmixing matrix $\mathbf{W}$ such that $\mathbf{Y}=\mathbf{W}\mathbf{X}$ is a good approximation of the independent components $\mathbf{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 non-Gaussianity (kurtosis or negentropy).

Figure: Figure 9: ICA estimated mixing matrix.
Refer to caption: /html/2410.11711/assets/x15.png

[Fig. 9](#A3.F9) shows the estimated mixing matrix $\mathbf{A}$ when running ICA on the D4RL-*expert* dataset on the Hopper environment. Under the assumptions of ICA, notably the statistical independence of the source signals, their linear mixing and the invertibility of the original (unknown) mixing matrix, the original sources are successfully recovered 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 scaling. In the case of our states and actions data, it’s not clear that this is the case from [Fig. 9](#A3.F9).
Similarly to PCA, we can transform the in-context multi-dimensional signal using ICA, and apply the ICL procedure to the recovered independent sources. We plan on exploring this method in future follow-up work.

### C.3 AutoEncoder-based approach

Variational Autoencoders (VAEs) (Kingma & Welling, ) offer a powerful framework for learning representations. A disentangled representation is one where each dimension of the latent space captures a distinct and interpretable factor of variation in the data. By combining an encoder network that maps inputs to a probabilistic latent space with a 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 $\beta$-VAE (Higgins et al., ) further emphasize this regularization by scaling the KL-divergence term, thereby encouraging the model to learn a more disentangled representation at the potential cost of reconstruction quality. Beyond simple VAEs, there exist previous work in the literature that specifically aim at learning a factorized posterior distribution in the latent space (Kim & Mnih, ). 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 a representation space, disregarding 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.

Figure: Figure 10: Sensitivity matrix.
Refer to caption: /html/2410.11711/assets/x16.png

#### Sensitivity Analysis.

Sensitivity analysis is a systematic approach to evaluate how the uncertainty in the output of a model can be attributed to different sources of uncertainty in the model’s inputs.
The One-at-a-Time (OAT) method is a technique used to understand the impact of individual input variables on the output of a model.
In the context of a transition function of a MDP, the OAT method involves systematically varying one current state or action dimension at a time, while keeping all others fixed, and observing the resulting changes in the output dimensions: $\frac{\partial(\mathbf{s}_{t+1})_{k}}{\partial(\mathbf{s}_{t})_{i}}\text{ and }\frac{\partial(\mathbf{s}_{t+1})_{k}}{\partial(\mathbf{a}_{t})_{j}}$, where $(\mathbf{s}_{t})_{i}$, $(\mathbf{a}_{t})_{j}$ and $(\mathbf{s}_{t+1})_{k}$ 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.

In practice, we measure the sensitivity by applying a perturbation (of scale $10\%$) to each input dimension separately, reporting the absolute change that occurs in each dimension of the output. Precisely, for a deterministic transition function $f$, input state dimension $i$, and output dimension $k$, we measure $|f(s+\epsilon,a)_{k}-f(s,a)_{k}|$ where $\epsilon_{i}=0.1\times scale(i)$ and 0 0 elsewhere.
The sensitivity matrix in [Fig. 10](#A3.F10) demonstrates that most of the next state dimensions are mostly affected by their respective previous values (the diagonal shape in the state dimensions square).
In addition to that, actions only directly affect some state dimensions, specifically velocities, which is expected from the nature of 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.1 .

This sensitivity analysis is specific to the single-step transition function. In practice, such conclusions might change when looking at a larger time scale of the simulation.

###### Remark C.1 .

## Appendix D Algorithms

### D.1 Soft-Actor Critic

Soft Actor-Critic (SAC) (Haarnoja et al., ) is an off-policy algorithm that incorporates the maximum entropy framework, which encourages exploration by seeking to maximize the entropy of the policy in addition to the expected return. SAC uses a deep neural network to approximate the policy (actor) and the value functions (critics), employing two Q-value functions to mitigate positive bias in the policy improvement step typical of off-policy algorithms. This approach helps in learning more stable and effective policies for complex environments, making SAC particularly suitable for tasks with high-dimensional, continuous action spaces.

We use the implementation provided in CleanRL (Huang et al., ) for SAC. In all environments, we keep the default hyperparameters provided with the library, except for the update frequency. In fact, the default update frequency in SAC is $1$ step, meaning that the policy that interacts with the environment gets updated after every interaction. In our LLM-based framework, this is problematic because it implies that the state visitation distribution of the in-context trajectories will be moving from one timestamp to another. To solve this issue, we fix the update frequency of SAC to the maximal number of steps of an episode of a given environment. With that in mind, we specify in [Table 1](#A4.T1) the complete list of hyperparameters used for every considered environment.

**Table 1: SAC hyperparameters.**
| Environment | HalfCheetah | Hopper | Pendulum |
| --- | --- | --- | --- |
| Update frequency | $1000$ | $1000$ | $200$ |
| Learning starts | $5000$ | $5000$ | $1000$ |
| Batch size | $128$ | $128$ | $64$ |
| Total timesteps | $1e6$ | $1e6$ | $1e4$ |
| Gamma $\gamma$ | $0.99$ | $0.99$ | $0.99$ |
| policy learning rate | $3\mathrm{e}-4$ | $3\mathrm{e}-4$ | $3\mathrm{e}-4$ |

### D.2 DICL-SAC

For our algorithm, we integrate an LLM inference interface (typically the Transformers library from Huggingface (Wolf et al., )) with CleanRL (Huang et al., ). [Table 2](#A4.T2) shows all DICL-SAC hyperparameter choices for the considered environments.

**Table 2: DICL-SAC hyperparameters.**
| Environment | HalfCheetah | Hopper | Pendulum |
| --- | --- | --- | --- |
| Update frequency | $1000$ | $1000$ | $200$ |
| Learning starts | $5000$ | $5000$ | $1000$ |
| LLM Learning starts | $10000$ | $10000$ | $2000$ |
| LLM Learning frequency | $256$ | $256$ | $16$ |
| Batch size | $128$ | $128$ | $64$ |
| LLM Batch size ($\alpha\%$) | $7(5\%),13(10\%),32(25\%)$ | $7(5\%),13(10\%),32(25\%)$ | $4(5\%),7(10\%),16(25\%)$ |
| Total timesteps | $1e6$ | $1e6$ | $1e4$ |
| Gamma $\gamma$ | $0.99$ | $0.99$ | $0.99$ |
| Max context length | $500$ | $500$ | $198$ |
| Min context length | $1$ | $1$ | $1$ |
| LLM sampling method | $mode$ | $mode$ | $mode$ |
| LLM dynamics learner | vICL | vICL | vICL |

#### Balancing gradient updates.

To ensure that DICL-SAC performs equally important gradient updates 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 $\mathcal{B}$ with the smallest batch size gets assigned a higher weight when doing gradient descent: $\frac{1}{|\mathcal{B}|}$. To address this, we multiply the loss corresponding to the LLM generated batch $\mathcal{B}_{\text{llm}}$ with a correcting coefficient $\frac{|\mathcal{B}_{\text{llm}}|}{|\mathcal{B}|}$ ensuring equal weighting across all samples.

We now show the full training curves on the HalfCheetah and Hopper environments ([Fig. 11](#A4.F11)). The return curves show smoothed average training curves $\pm$ 95% Gaussian confidence intervals for 5 seeds in HalfCheetah and Hopper, and 10 seeds for Pendulum.

Figure: Figure 11: Data-augmented off-policy RL. Full training curves.
Refer to caption: /html/2410.11711/assets/x17.png

## Appendix E What is the impact of the policy on the prediction error?

Figure: Figure 12: Correlation plots between state coverage and entropy of policies with MSE and KS metrics under the vICL dynamics learner.
Refer to caption: /html/2410.11711/assets/x20.png

In this experiment, We investigate how a policy impacts the accuracy and calibration of our LLM-based dynamics models. To do so, we train three model-free algorithms (PPO (Schulman et al., ), SAC (Haarnoja et al., ), and TD3 (Fujimoto et al., )) on the HalfCheetah environment, 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-Smirnov (KS) statistic. Our findings indicate that the state coverage 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 it also correlates with MSE, but interestingly, it does not appear to impact the calibration.

## Appendix F Multi-step prediction errors

#### The average multi-step error.

In [Fig. 3(a)](#S3.F3.sf1), we compute the average Mean Squared Error over prediction horizons for $h=1,\ldots,20$, and $5$ trajectories sampled uniformly from the D4RL expert dataset. For visualization purposes, we first rescale all the dimensions (using a pipeline composed of a *MinMaxScaler* and a *StandardScaler*) so that the respective MSEs are on the same scale.

#### The MLP baseline.

For the $MLP$ baseline, we instantiate an MLP with: $4$ layers, $128$ neurons each, and $ReLU$ activations. We then format the in-context trajectory as a dataset of $\{(s_{t},a_{t},s_{t+1})\}$ on which we train the MLP for $150$ epochs using early stopping and the Adam optimizer (Kingma & Ba, ).

We now extend [Fig. 3](#S3.F3) to show the multi-step generated trajectories for all the dimensions of the HalfCheetah system in [Fig. 13](#A6.F13).

Figure: Figure 13: Halfcheetah
Refer to caption: /html/2410.11711/assets/x21.png

## Appendix G Calibration

#### The naive baseline.

In the calibration plots [Figs. 7](#S4.F7) and [14](#A7.F14), we compare the LLM-based dynamics models with a (naive) baseline that estimates a Gaussian distribution using the in-context moments (mean and variance).

Kolmogorov-Smirnov Statistic (KS): This metric is computed using the quantiles (under the model distribution) of the ground truth values. Hypothetically, these quantiles are uniform if the error in predicting the ground truth is a random variable distributed according to a Gaussian with the predicted standard deviation, a property we characterize as *calibration*. To assess this, we compute the Kolmogorov-Smirnov (KS) statistics. Formally, starting from the model cumulative distribution function (CDF) $F_{\theta}(s_{t+1}|s_{t},a_{t})$, we define the empirical CDF of the quantiles of ground truth values by $\mathcal{F}_{\theta,j}(x)=\frac{\big{|}\big{\{}(s_{t},a_{t},s_{t+1})\in\mathcal{D}|F^{j}_{\theta}(s_{t+1}|s_{t},a_{t})\leq x\big{\}}\big{|}}{N}$ for $x\in[0,1]$. We denote by $U(x)$ the CDF of the uniform distribution over the interval $[0,1]$, and we define the KS statistics as the largest absolute difference between the two CDFs across the dataset $\mathcal{D}$:

$$ $\begin{split}\text{KS}(\mathcal{D};\theta;j\in\{1,\ldots,d_{s}\})&=\\ &\max_{i\in\{1,\ldots,N\}}\left|\mathcal{F}_{\theta,j}(F^{j}_{\theta}(s_{i,t+1}|s_{i,t},a_{i,t}))-U(F^{j}_{\theta}(s_{i,t+1}|s_{i,t},a_{i,t}))\right|\end{split}$ (10) $$

The KS score ranges between zero and one, with lower values indicating better calibration.

Figure: Figure 14: Halfcheetah
Refer to caption: /html/2410.11711/assets/x22.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Typo - on Line 63, should it be ‘deferred’ instead of ‘differed’?
2. Typo - on Line 322, should it be ‘The goal is to improve …’?
3. In the DICL-SAC algorithm, what would be the optimal value for \alpha? Are there any intuitions for choosing the optimal \alpha value?
soundness: 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.
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.
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.

This paper has very good potential, but its current form could benefit a lot from a systematic revision that improves its clarity and presentation.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 2
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?
soundness: 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.
2. 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.
3. 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.
4. It demonstrates that LLMs can act as reliable uncertainty estimators, a desirable trait for MBRL algorithms.
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.
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.
2. 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.
3. 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.
4. The writing style of the paper is a bit convoluted, making it less fluent to read.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: n/a
soundness: 3
strengths: (please also see the summary)
1. This work systematically studies LLM-based MBRL.  
2. The method is clearly presented and easy to follow.
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.

The 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.

Another 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.
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.

Another 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.

Finally, 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.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 3
strengths: 1. The authors proposed an innovative integration of LLMs into continuous reinforcement learning tasks.
2. Clear theoretical foundation supporting the DICL framework is shown.
3. The authors conduct extensive experiments to verify the effectiveness of the method across diverse RL scenarios.
4. The authors also did insightful analysis of zero-shot capabilities in complex MDPs.
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.
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.

Meta-review: This work explores using pre-trained LLMs to predict the dynamics of continuous MDPs in-context, addressing challenges related to multivariate data and control signal integration. The reviewers appreciate the originality of the approach, the theoretical analysis, and the thorough empirical evaluation. They also raise concerns, however, primarily on the clarity of presentation, on a few issues with evaluation, and limited discussion on generalizability. The authors respond to all concerns and questions in detail but not all reviewers engage in discussion.

Decision: Accept (Poster)