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

TL;DR

This work addresses zero-shot dynamics learning for continuous-state MDPs by leveraging pre-trained LLMs. It introduces Disentangled In-Context Learning (DICL), which projects trajectories into a disentangled space (via varphi, e.g., PCA), forecasts each component with an LLM in-context, and reconstructs the next-state trajectory with varphi^{-1}. The authors provide a multi-branch rollout return bound as a theoretical guarantee, and demonstrate practical use-cases including data-augmented offline RL with DICL-SAC and hybrid online/offline policy evaluation, all while showing well-calibrated uncertainty estimates through quantile calibration and reliability analyses. Empirically, DICL variants improve multi-step prediction accuracy and uncertainty calibration over baselines on proprioceptive tasks (e.g., HalfCheetah, Hopper) and offer meaningful sample-efficiency gains, albeit with trade-offs in branching and computation. The work advances the integration of LLMs into model-based RL by delivering a principled, zero-shot framework with theoretical guarantees and practical RL benefits, made reproducible by open-source code.

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.

Figures (20)

• Figure 1: The DICL Framework. DICL projects trajectories into a disentangled feature space before performing zero-shot forecasting using a pre-trained LLM and in-context learning.
• Figure 2: LLM can perceive time patterns. The LLM (Llama 3-8B) 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, focusing more on their first occurrence.
• Figure 3: The covariance matrix from an expert dataset in the Halfcheetah environment indicates linear correlations between state-action features.
• Figure 4: PCA-based DICL achieves smaller multi-step error in less computational time. We compare DICL-$(s)$ and DICL-$(s,a)$ using a number of components equal to half the number of features, with the vanilla approach vICL and an MLP baseline. (Llama 3-8B).
• Figure 5: 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, with the expectation over the overlapping branches as the return.

Theorems & Definitions (7)

• Definition 4.1: Multi-branch rollout return
• Theorem 4.2: Multi-branch return bound
• Lemma A.1
• proof
• Theorem A.2: Multi-branch return bound
• proof
• Remark C.1