Hierarchical Prompt Decision Transformer: Improving Few-Shot Policy Generalization with Global and Adaptive Guidance

TL;DR

HPDT addresses the challenge of few-shot generalization in Offline Meta-RL by introducing hierarchical prompting for Decision Transformers. It learns a global token to capture task-level transition dynamics and rewards and retrieves adaptive tokens from demonstrations to provide timestep-specific guidance, aided by Time2Vec time embeddings. Across seven MuJoCo and MetaWorld tasks, HPDT consistently outperforms baselines, including PDT variants and full fine-tuning, demonstrating strong in-context learning with improved data efficiency. The approach offers a lightweight, context-aware mechanism for adapting policies to unseen tasks without requiring extensive retraining, advancing practical cross-task generalization in offline RL.

Abstract

Decision transformers recast reinforcement learning as a conditional sequence generation problem, offering a simple but effective alternative to traditional value or policy-based methods. A recent key development in this area is the integration of prompting in decision transformers to facilitate few-shot policy generalization. However, current methods mainly use static prompt segments to guide rollouts, limiting their ability to provide context-specific guidance. Addressing this, we introduce a hierarchical prompting approach enabled by retrieval augmentation. Our method learns two layers of soft tokens as guiding prompts: (1) global tokens encapsulating task-level information about trajectories, and (2) adaptive tokens that deliver focused, timestep-specific instructions. The adaptive tokens are dynamically retrieved from a curated set of demonstration segments, ensuring context-aware guidance. Experiments across seven benchmark tasks in the MuJoCo and MetaWorld environments demonstrate the proposed approach consistently outperforms all baseline methods, suggesting that hierarchical prompting for decision transformers is an effective strategy to enable few-shot policy generalization.

Figures (4)

• Figure 1: The overall framework of the proposed HPDT. The top panel shows a demonstration trajectory from $\mathcal{D}_{\mathcal{M}_i}$ with a fictional blind lion looking for a red diamond. Blue circles describe different time stamps. The left part of the middle panel shows how we learn soft global token to capture the transition dynamic and reward function (by using the blue circles of the demonstration). The right part of the middle panel illustrates how to retrieve adaptive tokens for specific time stamps in rollout using KNN against the upper demonstration. The bottom panel shows how we fuse the global token (purple circle) and the adaptive tokens (blue circles) within a current rollout trajectory $\mathcal{O}_{\mathcal{M}_i}$.
• Figure 2: Meta-testing average performance of HPDT against baselines run over three random seeds facing unseen tasks. The $x$-axis represents the training epoch and $y$-axis is the average accumulated return on testing tasks.
• Figure 3: Ablation studies on Cheetah-Vel , Cheetah-Dir, and Pick&Place. In (a)(b)(c), we compare each ablation with the full model on Cheetah-Dir. Test tasks include running forward and backward. We show the accumulated reward for each task. The solid lines represent the full model HPDT for both tasks. The dashed lines represent the result of each ablation version. For Cheetah-Dir, the global token is more important. In (d)(e), we show the results for Cheetah-Vel and Pick&Place, where the adaptive tokens are more important. Curves represent the average accumulated reward on test tasks.
• Figure 4: 2D projections of the global tokens $g_{\mathcal{M}_i}^z$.