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

TL;DR

This paper introduces a hierarchical prompting method with retrieval augmentation for decision transformers, significantly enhancing few-shot policy generalization by providing context-specific guidance.

Contribution

It proposes a novel hierarchical prompting approach with global and adaptive tokens, improving few-shot reinforcement learning performance over static prompt methods.

Findings

Outperforms baseline methods on seven MuJoCo and MetaWorld tasks.

Demonstrates improved policy generalization with hierarchical prompts.

Adaptive tokens enable context-aware guidance in decision transformers.

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.