Research map / Large language models / Language, safety and interpretability
Zero-shot Model-based Reinforcement Learning using Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DICL, which PCA-disentangles state-action trajectories and feeds each component to a pre-trained LLM (Llama 3 family) via digit-token in-context forecasting, applied to model-based policy evaluation and SAC data augmentation.
Problem
Using pre-trained LLMs for zero-shot in-context prediction of continuous-state MDP dynamics is understudied; multivariate data and action integration are the key obstacles.
Approach
Disentangled In-Context Learning projects state-action trajectories with PCA into a linearly uncorrelated space, forecasts each component independently with an LLM via digit tokenization, and maps predictions back with the inverse transform; LLM logits give calibrated uncertainty.
Key results
- DICL-(s) with Llama 3.1 70B achieves the lowest multi-step MSE (3.59e-2) across LLM variants, vs. 4.06e-2 for the MLP baseline.
- DICL-(s,a) gives the best KS calibration statistic (45e-2 with Llama 3.1 8B) versus 55e-2 for MLP.
- DICL-SAC improves SAC sample efficiency early in training on Pendulum and HalfCheetah, with moderate gains.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply PCA to state features (DICL-(s)) instead of forecasting dimensions independently (vICL) | Llama-based dynamics forecasting · HalfCheetah and other control tasks | multi-step MSE | Helped |
| Apply PCA to state and action features (DICL-(s,a)) instead of vICL | LLM dynamics forecaster · 7 D4RL tasks | KS/10^-2: 53 $\pm$ 7 -> \textbf{45 $\pm$ 6} | Helped |
| Include action features in the PCA projection and forecasting | DICL-(s,a) · 7 D4RL tasks | MSE/10^-2 | Hurt |
| Use LLM-generated transitions to augment SAC replay data | SAC · Pendulum and HalfCheetah | return | Helped |
| Increase the proportion of LLM-generated data | DICL-SAC · HalfCheetah and other control environments | return | Hurt |
| Use DICL-(s) or DICL-(s,a) instead of vICL for reward-trajectory prediction | Hybrid policy evaluation · Hopper | relative value error | Helped |
| Use vICL instead of DICL for reward-trajectory prediction | Hybrid policy evaluation · HalfCheetah | relative value error | Helped |
| Vary the number of PCA components (including using only 4 of 23) (ablation) | DICL-(s,a) · HalfCheetah | average multi-step error | Mixed |
| Use PCA-based dimensionality reduction, reducing the original number of features by half | DICL · Control-task dynamics forecasting | computational time | Helped |
Surprising findings
- DICL-(s,a) has the highest prediction MSE across all LLMs despite better calibration, attributed to the extra error from predicting actions.
- In HalfCheetah, vanilla vICL outperforms DICL on reward-trajectory prediction for policy evaluation.
Limitations
- Only LLaMA-series models are used because their tokenizer assigns single tokens to numbers 0-999; other LLMs need logit-extraction schemes.
- Predicting actions with DICL-(s,a) autoregressively is computationally expensive for model-based RL, as the paper acknowledges.
- PCA is linear and may not capture nonlinear dependencies.
Experimental setup
Datasets: D4RL MuJoCo HalfCheetah Hopper Pendulum
Benchmarks: D4RL HalfCheetah Hopper Pendulum
Metrics: MSE Kolmogorov-Smirnov statistic relative value error quantile calibration
Baselines: vICL MLP SAC
Builds on: LLMTime Model-Based Policy Optimization Soft Actor-Critic PCA
Closest papers by meaning
- Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts (NeurIPS 2025 · similarity 0.90)
- Efficient Reinforcement Learning with Large Language Model Priors (ICLR 2025 · similarity 0.89)
- XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.89)
- Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.89)
- ShiQ: Bringing back Bellman to LLMs (NeurIPS 2025 · similarity 0.88)
- Can large language models explore in-context? (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.