Pre-trained Language Models Improve the Few-shot Prompt Ability of Decision Transformer

TL;DR

This paper introduces LPDT, a framework that uses pretrained language models and low-rank adaptation to improve few-shot decision transformer performance in offline reinforcement learning, reducing data requirements.

Contribution

The paper proposes LPDT, combining pretrained language models with LoRA and prompt regularization, to enhance task differentiation and performance in meta-RL with limited data.

Findings

LPDT achieves comparable performance to Prompt-DT with only 10% of the data.

Pretrained language models provide rich prior knowledge beneficial for RL tasks.

Ablation studies confirm the effectiveness of each component in LPDT.

Abstract

Decision Transformer (DT) has emerged as a promising class of algorithms in offline reinforcement learning (RL) tasks, leveraging pre-collected datasets and Transformer's capability to model long sequences. Recent works have demonstrated that using parts of trajectories from training tasks as prompts in DT enhances its performance on unseen tasks, giving rise to Prompt-DT methods. However, collecting data from specific environments can be both costly and unsafe in many scenarios, leading to suboptimal performance and limited few-shot prompt abilities due to the data-hungry nature of Transformer-based models. Additionally, the limited datasets used in pre-training make it challenging for Prompt-DT type of methods to distinguish between various RL tasks through prompts alone. To address these challenges, we introduce the Language model-initialized Prompt Decision Transformer (LPDT) framework, which leverages pretrained language models providing rich prior knowledge for RL tasks and fine-tunes the sequence model using Low-rank Adaptation (LoRA) for meta-RL problems. We further incorporate prompt regularization to effectively differentiate between tasks based on prompt feature representations. Comprehensive empirical studies demonstrate that initializing with a pre-trained language model provides the prior knowledge and achieves a similar performance with Prompt-DT under only $10\%$ data in some MuJoCo control tasks. We also provide a thorough ablation study to validate the effectiveness of each component, including sequence modeling, language models, prompt regularizations, and prompt strategies.