Prompting Decision Transformer for Few-Shot Policy Generalization

Mengdi XuYikang ShenShun ZhangYuchen LuDing ZhaoJoshua B. TenenbaumChuang Gan

article2022ICML214 citations

Presents Prompt-DT, a sequence-modeling approach that conditions decision transformers on short demonstration trajectory prompts to enable offline policy adaptation to unseen and out-of-distribution control tasks without fine-tuning.

Listen

Deploying reinforcement learning agents in high-stakes, real-world domains such as robotics, healthcare, and autonomous driving requires learning from pre-collected historical data rather than dangerous or costly live trial-and-error. However, conventional offline reinforcement learning methods struggle significantly when faced with new, unseen tasks, requiring complex algorithms and computationally heavy model updates to adapt.

The article demonstrates that conditioning an autoregressive Transformer architecture on short demonstration segments—termed trajectory prompts—enables an artificial intelligence agent to rapidly generalize to unseen control tasks without any test-time model parameter updates or fine-tuning.

The authors evaluate their proposed method, Prompt-based Decision Transformer (Prompt-DT), across five continuous control benchmarks in simulated robotic environments. The approach treats decision-making as a sequence-modeling problem, prepending very short demonstration trajectories (ranging from 2 to 15 timesteps) directly to the agent's recent history to specify the target task. Prompt-DT is compared against traditional multi-task models and established meta-learning algorithms, with performance measured by cumulative task rewards across multiple test environments.

The evaluation yields several critical findings. First, Prompt-DT consistently outperforms strong offline meta-reinforcement learning baselines such as MACAW by a wide margin, converging faster and achieving higher performance while requiring substantially less adaptation data (for instance, 5 timesteps versus 256 samples). Second, trajectory prompts provide far more effective task identification than historical reward signals alone, allowing the agent to adapt immediately without test-time fine-tuning. Third, the system exhibits strong robustness to prompt length—achieving near-peak performance with as few as two timesteps—but displays high sensitivity to prompt quality, performing best when conditioned on high-quality demonstrations. Finally, Prompt-DT successfully generalizes to out-of-distribution tasks that lie beyond the range of training goals, an evaluation setting where prior methods fail.

These findings suggest that inductive architecture design and prompting techniques can replace complex meta-optimization loops, reducing the computational cost, training time, and engineering overhead of deploying adaptable autonomous agents. By eliminating online fine-tuning, this framework minimizes the risk of catastrophic forgetting and catastrophic failures during field deployment.

Organizations developing autonomous systems should explore prompt-conditioned sequence architectures over gradient-based meta-learning when rapid task switching is required from pre-collected data. When implementing these systems, teams should prioritize the collection of high-quality demonstration prompts rather than large volumes of adaptation data. Before deploying in production, further research is recommended to address current limitations, particularly the model's struggle on highly complex compositional tasks (such as multi-stage tool manipulation in the ML10 benchmark) and its dependence on high-quality offline datasets.

Carnegie Mellon University Massachusetts Institute of Technology Mila MIT-IBM Watson AI Lab Université de Montréal University of Massachusetts Amherst encoder-decoder Few-Shot Policy Generalization Decision Transformer Prompt-DT Architecture

arXiv: 2206.13499 mxu34/prompt-dt

Paper: Decision Transformer: Reinforcement Learning via Sequence Modeling, Lili Chen et al. (2021). It introduces the Decision Transformer architecture that reformulates offline reinforcement learning as conditional sequence modeling, establishing the foundational policy backbone that Prompt-DT builds upon.
Paper: Making Pre-trained Language Models Better Few-shot Learners, Tianyu Gao et al. (2021). It provides essential background on few-shot adaptation using demonstration prompts and prompt-based tuning that motivates Prompt-DT's trajectory prompt framework.
Paper: Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing, Pengfei Liu et al. (2021). It presents a comprehensive taxonomy and theoretical foundation of prompting paradigms across sequence models, clarifying the prompting mechanics adapted for policy generalization.
Paper: The Power of Scale for Parameter-Efficient Prompt Tuning, Brian Lester et al. (2021). It investigates parameter-efficient soft prompt conditioning in Transformers, providing architectural insight into how prompt tokens steer sequence models without full-parameter retraining.
Paper: Offline Reinforcement Learning with Implicit Q-Learning, Ilya Kostrikov et al. (2021). It details key concepts and benchmark standards in offline reinforcement learning without online environmental interactions, framing the offline policy optimization paradigm used in Prompt-DT.

Paper: Supervised Pretraining Can Learn In-Context Reinforcement Learning, Jonathan Lee et al. (2023). It generalizes in-context reinforcement learning via supervised pretraining of Transformers across online and offline tasks, extending Prompt-DT's prompt-guided policy formulation.
Paper: Multi-Game Decision Transformers, Kuang-Huei Lee et al. (2022). It scales the decision-transformer sequence modeling approach across dozens of multi-game environments, expanding on the trajectory-conditioned generalist agent paradigm.
Paper: Human-Timescale Adaptation in an Open-Ended Task Space, Jakob Bauer et al. (2023). It explores large-scale fast adaptation and in-context policy improvement in complex 3D environments, continuing Prompt-DT's investigation into architecture-driven few-shot RL.
Paper: PIVOT: Iterative Visual Prompting Elicits Actionable Knowledge for VLMs, Soroush Nasiriany et al. (2024). It applies iterative prompting to vision-language models for zero-shot and few-shot continuous robotic control, building on prompt-based decision-making concepts.
Paper: A Survey on In-context Learning, Qingxiu Dong et al. (2024). It synthesizes emerging theories and empirical mechanisms of in-context learning and demonstration conditioning, contextualizing trajectory-based prompting methods.

Cover for Prompting Decision Transformer for Few-Shot Policy Generalization

Abstract

Human can leverage prior experience and learn novel tasks from a handful of demonstrations. In contrast to offline meta-reinforcement learning, which aims to achieve quick adaptation through better algorithm design, we investigate the effect of architecture inductive bias on the few-shot learning capability. We propose a Prompt-based Decision Transformer (Prompt-DT), which leverages the sequential modeling ability of the Transformer architecture and the prompt framework to achieve few-shot adaptation in offline RL. We design the trajectory prompt, which contains segments of the few-shot demonstrations, and encodes task-specific information to guide policy generation. Our experiments in five MuJoCo control benchmarks show that Prompt-DT is a strong few-shot learner without any extra finetuning on unseen target tasks. Prompt-DT outperforms its variants and strong meta offline RL baselines by a large margin with a trajectory prompt containing only a few timesteps. Prompt-DT is also robust to prompt length changes and can generalize to out-of-distribution (OOD) environments. Project page: https://mxu34.github.io/PromptDT/.

Topics

encoder-decoder

Keywords

Few-Shot Policy Generalization Decision Transformer Prompt-DT Architecture meta-reinforcement learning stochastic trajectory prompt Model-Agnostic Meta-Learning sequential decision-making task Meta-Actor Critic Robot Learning

1. Introduction
2. Related Work
3. Preliminaries
3.1. Online and Offline Meta-Reinforcement Learning
3.2. Decision Transformer
4. Prompt-based Decision Transformer
4.1. Problem Formulation
4.2. Prompt Representation
4.3. Prompt-DT Architecture
4.4. Algorithms
5. Experiments
5.1. Environments and Datasets
5.2. Baselines
6. Discussion
6.1. Can Prompt-DT Achieve Few-Shot Policy Generalization?
6.2. Does the Prompt Quantity Affect the Few-Shot Generalization Ability?
6.3. Does the Prompt Quality Affect the Few-Shot Generalization Ability?
6.4. Can Prompt-DT Generalize to Out-of-distribution Tasks via Few-Shot Demonstrations?
7. Conclusion
References
Appendix: Prompting Decision Transformer for Few-Shot Policy Generalization
A. Hyperparameters
B. Experiment Details
C. Ablation Study
C.1. The Effect of Prompt Quantity
C.2. The Effect of Finetune Data's Quantity on MT-BC-Finetune
C.3. The Effect of Finetune Data's Quantity on MACAW

Knowls

Knowl 1 — Prompt-based Decision Transformer (Prompt-DT) Architecture
model/method
The Prompt-based Decision Transformer (Prompt-DT) is an architecture designed for offline few-shot policy generalization. Prompt-DT casts policy generation as a conditional sequence modeling task using a causal GPT-style Transformer.

Given a task $\mathcal{T}_i$ , Prompt-DT takes an input sequence $\tau^{\text{input}} = (\tau_i^\star, \tau_i)$ , where $\tau_i^\star$ is a trajectory prompt containing $K^\star$ timesteps sampled from a demonstration dataset $\mathcal{P}_i$ , and $\tau_i$ is the agent's recent context history containing $K$ timesteps sampled from the task's dataset $\mathcal{D}_i$ .

Each timestep in both the prompt and the history is represented as a 3-tuple of reward-to-go $\hat{r}$ , state $s$ , and action $a$ . At timestep $t$ , the reward-to-go is $\hat{r}_t = \sum_{t'=t}^T r_{t'}$ . Consequently, the full input sequence $\tau^{\text{input}}$ corresponds to $3(K^\star + K)$ tokens:

$\tau^{\text{input}} = (\hat{r}_1^\star, s_1^\star, a_1^\star, \dots, \hat{r}_{K^\star}^\star, s_{K^\star}^\star, a_{K^\star}^\star, \hat{r}_{t-K+1}, s_{t-K+1}, a_{t-K+1}, \dots, \hat{r}_t, s_t, a_t)$

Token embeddings are produced via linear projections of each modality ( $s, a, \hat{r}$ ) added to a learned positional embedding corresponding to the timestep. A causal attention mask is applied so that tokens only attend to preceding tokens. Prompt-DT autoregressively predicts actions at the heads corresponding to state tokens by optimizing the mean-squared error loss across all $K^\star + K$ action positions.
Knowl 2 — Trajectory Prompt Representation in Offline RL
definition
A trajectory prompt $\tau_i^\star$ for a Markov Decision Process (MDP) task $\mathcal{T}_i = (\mathcal{S}, \mathcal{A}, P_i, R_i, \mu)$ is a sequence of state, action, and reward-to-go transitions sampled from a few-shot demonstration dataset $\mathcal{P}_i$ . It implicitly encapsulates the transition dynamics $P_i$ and reward structure $R_i$ of the specific task without requiring language templates or textual annotations.

To ensure training stability and prevent overfitting to a single trajectory, the trajectory prompt $\tau_i^\star$ is defined stochastically as an aggregation of $J$ trajectory segments, each of length $H$ , yielding a total prompt length of $K^\star = JH$ steps:

$\tau_i^\star = (\tau_{i,1}^\star, \dots, \tau_{i,J}^\star)$

where each segment $\tau_{i,j}^\star$ is defined as:

$\tau_{i,j}^\star = (\hat{r}_{i,j,1}^\star, s_{i,j,1}^\star, a_{i,j,1}^\star, \dots, \hat{r}_{i,j,H}^\star, s_{i,j,H}^\star, a_{i,j,H}^\star), \quad \forall j \in \{1, \dots, J\}$

Here, $s_{i,j,h}^\star \in \mathcal{S}$ is the state, $a_{i,j,h}^\star \in \mathcal{A}$ is the action, and $\hat{r}_{i,j,h}^\star = \sum_{t'=h}^{H} r_{i,j,t'}^\star$ is the segment reward-to-go. The prompt horizon $K^\star$ is set significantly shorter than the full task horizon $T$ , providing sufficient task-identification signal while remaining concise.
Knowl 3 — Prompt-DT Training and Prompt Generation Algorithms
algorithm
Prompt-DT is trained offline over a distribution of training tasks $\mathcal{T}^{\text{train}}$ by minimizing action prediction errors over both prompt tokens and trajectory tokens, stabilizing multi-task training by aggregating gradient estimates over batches from all training tasks simultaneously.
Algorithm: Trajectory Prompt Generation (GetPrompt)
Input: Task $\mathcal{T}$ , demonstration dataset $\mathcal{P}$ , segment count $J$ , segment length $H$
Sample $J$ episodes from $\mathcal{P}$
for $j = 1$ to $J$ do
    Sample a continuous trajectory segment $\tau_j^\star = (\hat{r}_{j,1}^\star, s_{j,1}^\star, a_{j,1}^\star, \dots, \hat{r}_{j,H}^\star, s_{j,H}^\star, a_{j,H}^\star)$ from episode $j$
end for
return Trajectory prompt $\tau^\star = (\tau_1^\star, \dots, \tau_J^\star)$
Algorithm: Prompt-DT Training
Input: Training tasks $\mathcal{T}^{\text{train}}$ , causal Transformer $\text{Transformer}_\theta$ , iterations $N$ , offline dataset $\mathcal{D}$ , demonstration set $\mathcal{P}$ , task batch size $M$ , learning rate $\alpha$
for $n = 1$ to $N$ do
    for each task $\mathcal{T}_i \in \mathcal{T}^{\text{train}}$ do
        for $m = 1$ to $M$ do
            Sample trajectory $\tau_{i,m}$ of length $K$ from $\mathcal{D}_i$
            Sample prompt $\tau_{i,m}^\star = \text{GetPrompt}(\mathcal{T}_i, \mathcal{P}_i)$
            Construct $\tau_{i,m}^{\text{input}} = (\tau_{i,m}^\star, \tau_{i,m})$
        end for
        Form task minibatch $\mathcal{B}_i^M = \{\tau_{i,m}^{\text{input}}\}_{m=1}^M$
    end for
    Form batch $\mathcal{B} = \bigcup_{\mathcal{T}_i \in \mathcal{T}^{\text{train}}} \mathcal{B}_i^M$
     $a^{\text{pred}} = \text{Transformer}_\theta(\tau^{\text{input}}), \quad \forall \tau^{\text{input}} \in \mathcal{B}$
     $\mathcal{L}_{\text{MSE}} = \frac{1}{|\mathcal{B}|} \sum_{\tau^{\text{input}} \in \mathcal{B}} \|a - a^{\text{pred}}\|^2$
     $\theta \leftarrow \theta - \alpha \nabla_\theta \mathcal{L}_{\text{MSE}}$
end for
Knowl 4 — Prompt-DT Zero-Gradient Few-Shot Evaluation Algorithm
algorithm
During evaluation on unseen test tasks $\mathcal{T}^{\text{test}}$ , Prompt-DT generates actions online by conditioning on a pre-sampled prompt and a streaming history buffer without fine-tuning model parameters $\theta$ or computing gradient updates.
Algorithm: Prompt-DT Few-Shot Evaluation
Input: Test tasks $\mathcal{T}^{\text{test}}$ , causal Transformer $\text{Transformer}_\theta$ , test demonstration set $\mathcal{P}$ , target returns $\{G_i^\star\}$ , episode length $T$
for each task $\mathcal{T}_i \in \mathcal{T}^{\text{test}}$ do
    Initialize history $\tau$ as all-zero placeholders of context length $K$
    Initialize remaining target return $g = G_i^\star$
    Sample prompt $\tau^\star = \text{GetPrompt}(\mathcal{T}_i, \mathcal{P}_i)$
    for $t = 1$ to $T$ do
        Get action $a_t = \text{Transformer}_\theta((\tau^\star, \tau))[-1]$
        Execute action $a_t$ in environment; observe state $s_{t+1}$ and reward $r_t$
        Update desired return-to-go: $g \leftarrow g - r_t$
        Append transition tuple $[s_{t+1}, a_t, g]$ to recent history $\tau$ , maintaining context length $K$
    end for
end for
Knowl 5 — Benchmark Setup and Baseline Configurations for Few-Shot Generalization
experimental setup
The few-shot policy generalization framework evaluates offline policy learning across five continuous control domains:
- Cheetah-dir: 2 tasks (forward, backward running directions), rewarded on directional velocity.
- Cheetah-vel: 40 tasks with target velocities uniformly sampled from $[0, 3]$ . 35 tasks are used for training and 5 held-out tasks for testing.
- Ant-dir: 50 tasks with 2D goal directions uniformly sampled on the circle. 45 tasks are used for training and 5 held-out tasks for testing.
- Dial: A 6-DOF Jaco robot reaching target numbers on a number pad (10 tasks total: 6 training, 4 testing).
- Meta-World reach-v2: A Sawyer robot end-effector reaching target 3D positions (20 tasks total: 15 training, 5 testing).
Prompt-DT is compared against four baselines under identical training data:
1. Multi-task Offline RL (MT-ORL): Standard Decision Transformer trained on multi-task data with reward-to-go conditioning but without trajectory prompts.
2. Prompt-based Behavior Cloning (Prompt-MT-BC): Prompt-DT variant where reward-to-go tokens $\hat{r}$ are omitted from the history input in both training and testing, conditioning only on $(s, a)$ history and prompt.
3. Multi-task Behavior Cloning with Finetuning (MT-BC-Finetune): Standard Decision Transformer without prompts or reward-to-go tokens, adapted to test tasks using gradient finetuning steps on target task data.
4. Meta-Actor Critic with Advantage Weighting (MACAW): An offline meta-RL baseline combining advantage-weighted regression with MAML-style gradient adaptation.
Knowl 6 — Few-Shot Generalization Performance Across MuJoCo Benchmarks
empirical result
Across never-before-seen test tasks in Cheetah-dir, Cheetah-vel, Ant-dir, Dial, and Meta-World reach-v2, Prompt-DT and Prompt-MT-BC consistently outperform non-prompt baselines (MT-ORL, MT-BC-Finetune, and MACAW) using short prompts ( $K^\star = 5$ for Cheetah-dir, Cheetah-vel, and Ant-dir; $K^\star = 2$ for Meta-World reach-v2; $K^\star = 15$ for Dial):
- Prompt vs. Return Conditioning (MT-ORL): Prompt-DT and Prompt-MT-BC outperform MT-ORL by large margins across environments. While reward-to-go tokens in context history fail to guide action selection on novel tasks, trajectory prompts provide explicit task identification signals.
- Prompt-DT vs. Prompt-MT-BC: In Cheetah-dir, Cheetah-vel, Ant-dir, and Meta-World reach-v2, Prompt-DT and Prompt-MT-BC achieve nearly identical returns, indicating the trajectory prompt itself provides sufficient task information. In Dial, Prompt-DT outperforms Prompt-MT-BC, indicating that reward-to-go information assists in tasks where short prompts alone leave task ambiguity.
- Prompt Adaptation vs. Gradient Adaptation (MACAW and MT-BC-Finetune): Prompt-based methods converge to optimal test return substantially faster than MACAW and MT-BC-Finetune during pretraining. Even when MACAW is provided with 256 adaptation samples and 10 gradient steps at test time, Prompt-DT with only 5 demonstration steps matches MACAW in Cheetah-dir and outperforms MACAW in Cheetah-vel and Ant-dir.

Knowl 7 — Effect of Trajectory Prompt Length and Segmentation on Generalization

data/table

The effect of the trajectory prompt length $K^\star = JH$ on few-shot generalization performance was evaluated across different numbers of segments $J$ and segment lengths $H$ , using context history length $K = 20$ over 3 random seeds.

$K^\star$	$J$	$H$	Cheetah-dir	Cheetah-vel	Ant-dir
2	1	2	$926.46 \pm 2.87$	$-45.26 \pm 2.87$	$409.81 \pm 9.69$
5	1	5	$927.20 \pm 18.02$	$-37.92 \pm 4.56$	$367.12 \pm 10.50$
10	1	10	$925.00 \pm 1.41$	$-38.43 \pm 2.14$	$382.94 \pm 25.21$
40	2	20	$926.87 \pm 9.30$	$-34.43 \pm 2.33$	$323.83 \pm 9.33$

The results show that Prompt-DT is robust to changes in prompt length:

In Cheetah-dir, Prompt-DT achieves maximal returns ( $approx 926$ ) with as few as $K^\star = 2$ timesteps, as the shared locomotion skill is maintained in history $\tau$ while direction is identified from single $(s, a, \hat{r})$ transitions.
In Cheetah-vel, performance improves slightly from $-45.26$ at $K^\star=2$ to $-34.43$ at $K^\star=40$ , but short prompts ( $K^\star=2$ ) remain competitive.
In Ant-dir, the highest return ( $409.81$ ) is achieved at $K^\star = 2$ , as multi-episode prompts risk incorporating sub-optimal transitions from non-expert trajectories present in the dataset.

Knowl 8 — Sensitivity of Prompt-DT to Demonstration and Training Quality
empirical result
The performance of Prompt-DT is heavily governed by the quality of the trajectory prompt and the coverage of the training dataset:
1. Expert Training Dataset: When trained on expert offline data in Cheetah-vel, Prompt-DT conditioned on expert prompts achieves near-optimal episode returns (approximately $-35$ ), conditioned on medium prompts maintains competitive performance (approximately $-50$ ), but conditioned on random prompts drops to poor returns (approximately $-200$ ). Prompt-DT modulates its generated policy directly according to the quality reflected in the prompt.
2. Medium Training Dataset: When trained on medium-quality data, Prompt-DT performance drops significantly with random prompts and marginally increases when prompted with expert demonstrations.
3. Random Training Dataset: When trained exclusively on random exploration data, Prompt-DT fails to achieve high returns even when conditioned on expert trajectory prompts (returns stagnate around $-250$ ), because the random training dataset's state distribution does not overlap with expert target states, causing the model to collapse to the common random state distribution.
Knowl 9 — Out-of-Distribution Policy Extrapolation in Ant-dir
empirical result
Prompt-DT exhibits out-of-distribution (OOD) policy extrapolation when evaluated on goals outside the numerical range seen during training.

In Ant-dir, 8 training tasks with direction angles restricted to a specific interval (task indices $\{8, 13, 16, 20, 22, 26, 32, 37\}$ ) were used for pretraining, and 3 test tasks with directions strictly outside the training interval (task indices $\{1, 4, 41\}$ ) were evaluated.

Prompt-augmented models (Prompt-DT and Prompt-MT-BC) achieved episodic returns exceeding $300$ to $500$ on OOD target directions, substantially outperforming non-prompt baselines (MT-ORL, MT-BC-Finetune, and MACAW, which stayed near $150$ to $200$ ). This demonstrates that trajectory prompts allow Transformer models to extrapolate target directions beyond the training support without parameter fine-tuning.
Knowl 10 — Limitations in Compositional Benchmarks and Prompt Weighting
limitation
Prompt-DT has several identified limitations:
1. Failure on Compositional Multi-Task Benchmarks: When evaluated on Meta-World's ML10 benchmark using prompts subsampled from expert trajectories, Prompt-DT and Prompt-MT-BC fail to generalize to novel compositional manipulation tasks, indicating that simple trajectory segments are insufficient when complex tool use and multi-stage sub-goals are required.
2. Fixed Prompt Hyperparameter: The prompt length $K^\star$ is treated as a fixed hyperparameter rather than dynamically adjusted or weighted relative to the history context $\tau$ .
3. Sensitivity to Low-Quality Replay Data: Prompt-DT requires high-quality multi-task pretraining data; pretraining purely on random exploration buffers prevents the model from leveraging expert test prompts.

Coverage note — Deliberately omitted general background discussions on NLP prompt templates, standard Decision Transformer background derivation, standard MAML derivation, and basic MuJoCo physics simulation details, as they represent established prior work rather than the contributed methods and findings of this paper.

References

1.Altae-Tran, H., Ramsundar, B., Pappu, A. S., and Pande, V. Low data drug discovery with one-shot learning. ACS central science, 3(4):283–293, 2017.
2.Bertinetto, L., Henriques, J. F., Valmadre, J., Torr, P., and Vedaldi, A. Learning feed-forward one-shot learners. In Advances in neural information processing systems, pp. 523–531, 2016.
3.Brock, A., Lim, T., Ritchie, J. M., and Weston, N. Smash: one-shot model architecture search through hypernetworks. arXiv preprint arXiv:1708.05344, 2017.
4.Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.
5.Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A., and Zagoruyko, S. End-to-end object detection with transformers. In European Conference on Computer Vision, pp. 213–229. Springer, 2020.
6.Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345, 2021.
7.Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
8.Dong, X., Zhu, L., Zhang, D., Yang, Y., and Wu, F. Fast parameter adaptation for few-shot image captioning and visual question answering. In Proceedings of the 26th ACM international conference on Multimedia, pp. 54–62, 2018.
9.Duan, Y., Andrychowicz, M., Stadie, B. C., Ho, J., Schneider, J., Sutskever, I., Abbeel, P., and Zaremba, W. One-shot imitation learning. arXiv preprint arXiv:1703.07326, 2017.
10.Ebert, F., Finn, C., Dasari, S., Xie, A., Lee, A., and Levine, S. Visual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint arXiv:1812.00568, 2018.
11.Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, pp. 1126–1135. PMLR, 2017a.
12.Finn, C., Yu, T., Zhang, T., Abbeel, P., and Levine, S. One-shot visual imitation learning via meta-learning. In Conference on Robot Learning, pp. 357–368. PMLR, 2017b.
13.Fujimoto, S., Hoof, H., and Meger, D. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, pp. 1587–1596. PMLR, 2018.
14.Furuta, H., Matsuo, Y., and Gu, S. S. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10364, 2021.
15.Gao, T., Fisch, A., and Chen, D. Making pre-trained language models better few-shot learners. arXiv preprint arXiv:2012.15723, 2020.
16.Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861–1870. PMLR, 2018.
17.Islam, R., Teru, K. K., Sharma, D., and Pineau, J. Off-Policy Policy Gradient Algorithms by Constraining the State Distribution Shift. arXiv:1911.06970 [cs, stat], 2019.
18.Janner, M., Li, Q., and Levine, S. Reinforcement learning as one big sequence modeling problem. arXiv preprint arXiv:2106.02039, 2021.
19.Kahn, G., Abbeel, P., and Levine, S. Badgr: An autonomous self-supervised learning-based navigation system. IEEE Robotics and Automation Letters, 6(2):1312–1319, 2021.
20.Kalashnikov, D., Irpan, A., Pastor, P., Ibarz, J., Herzog, A., Jang, E., Quillen, D., Holly, E., Kalakrishnan, M., Vanhoucke, V., et al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning, pp. 651–673. PMLR, 2018.
21.Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. MOReL : Model-Based Offline Reinforcement Learning. In arXiv:2005.05951 [Cs, Stat], 2021.
22.Kumar, A., Zhou, A., Tucker, G., and Levine, S. Conservative Q-Learning for Offline Reinforcement Learning. Neural Information Processing Systems, 2020.
23.Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
24.Liu, P., Yuan, W., Fu, J., Jiang, Z., Hayashi, H., and Neubig, G. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. arXiv preprint arXiv:2107.13586, 2021.
25.Mitchell, E., Rafailov, R., Peng, X. B., Levine, S., and Finn, C. Offline meta-reinforcement learning with advantage weighting. In International Conference on Machine Learning, pp. 7780–7791. PMLR, 2021.
26.Nichol, A., Achiam, J., and Schulman, J. On First-Order Meta-Learning Algorithms. arXiv:1803.02999 [cs], 2018.
27.Peng, X. B., Kumar, A., Zhang, G., and Levine, S. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177, 2019.
28.Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
29.Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. Meta-Learning with Implicit Gradients. In Wallach, H., Larochelle, H., Beygelzimer, A., d\textquotesingle Alche-Buc, F., Fox, E., and Garnett, R. (eds.), ´ Advances in Neural Information Processing Systems 32, pp. 113–124. Curran Associates, Inc., 2019.
30.Rothfuss, J., Lee, D., Clavera, I., Asfour, T., and Abbeel, P. Promp: Proximal meta-policy search. arXiv preprint arXiv:1810.06784, 2018.
31.Shiarlis, K., Wulfmeier, M., Salter, S., Whiteson, S., and Posner, I. Taco: Learning task decomposition via temporal alignment for control. In International Conference on Machine Learning, pp. 4654–4663. PMLR, 2018.
32.Shridhar, M., Thomason, J., Gordon, D., Bisk, Y., Han, W., Mottaghi, R., Zettlemoyer, L., and Fox, D. Alfred: A benchmark for interpreting grounded instructions for everyday tasks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10740–10749, 2020.
33.Sutton, R. S. and Barto, A. G. Reinforcement Learning: An Introduction. MIT Press, 2018. ISBN 978-0-262-35270-3.
34.Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026–5033. IEEE, 2012.
35.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In Advances in neural information processing systems, pp. 5998–6008, 2017.
36.Vinyals, O., Blundell, C., Lillicrap, T., Wierstra, D., et al. Matching networks for one shot learning. Advances in neural information processing systems, 29:3630–3638, 2016.
37.Wang, Y., Yao, Q., Kwok, J. T., and Ni, L. M. Generalizing from a few examples: A survey on few-shot learning. ACM Computing Surveys (CSUR), 53(3):1–34, 2020.
38.Wu, Y. and Demiris, Y. Towards one shot learning by imitation for humanoid robots. In 2010 IEEE International Conference on Robotics and Automation, pp. 2889–2894. IEEE, 2010.
39.Yoon, J., Kim, T., Dia, O., Kim, S., Bengio, Y., and Ahn, S. Bayesian model-agnostic meta-learning. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 7343–7353, 2018.
40.Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning, pp. 1094–1100. PMLR, 2020a.
41.Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J., Levine, S., Finn, C., and Ma, T. MOPO: Model-based Offline Policy Optimization. Neural Information Processing Systems, 2020b.
42.Yu, T., Kumar, A., Rafailov, R., Rajeswaran, A., Levine, S., and Finn, C. COMBO: Conservative Offline Model-Based Policy Optimization. arXiv:2102.08363 [cs], 2021.

Citation

MLA

Xu, M., et al. “Prompting Decision Transformer for Few-Shot Policy Generalization”. International Conference on Machine Learning, vol. 162, 2022, pp. 24631–45, https://proceedings.mlr.press/v162/xu22g.html.

APA

Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., & Gan, C. (2022). Prompting Decision Transformer for Few-Shot Policy Generalization. International Conference on Machine Learning, 162, 24631–24645. https://proceedings.mlr.press/v162/xu22g.html

Chicago

Xu, M., Y. Shen, S. Zhang, et al. 2022. “Prompting Decision Transformer for Few-Shot Policy Generalization”. International Conference on Machine Learning 162: 24631–45. https://proceedings.mlr.press/v162/xu22g.html.

Harvard

Xu, M. et al. (2022) “Prompting Decision Transformer for Few-Shot Policy Generalization”, International Conference on Machine Learning. PMLR, pp. 24631–24645. Available at: https://proceedings.mlr.press/v162/xu22g.html.

Vancouver

1. Xu M, Shen Y, Zhang S, Lu Y, Zhao D, Tenenbaum J, Gan C (2022) Prompting Decision Transformer for Few-Shot Policy Generalization. In: International Conference on Machine Learning. PMLR, pp 24631–24645

BibTeX

@InProceedings{pmlr-v162-xu22g,
  title = 	 {Prompting Decision Transformer for Few-Shot Policy Generalization},
  author =       {Xu, Mengdi and Shen, Yikang and Zhang, Shun and Lu, Yuchen and Zhao, Ding and Tenenbaum, Joshua and Gan, Chuang},
  booktitle = 	 {Proceedings of the 39th International Conference on Machine Learning},
  pages = 	 {24631--24645},
  year = 	 {2022},
  editor = 	 {Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan},
  volume = 	 {162},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {17--23 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v162/xu22g/xu22g.pdf},
  url = 	 {https://proceedings.mlr.press/v162/xu22g.html},
  abstract = 	 {Human can leverage prior experience and learn novel tasks from a handful of demonstrations. In contrast to offline meta-reinforcement learning, which aims to achieve quick adaptation through better algorithm design, we investigate the effect of architecture inductive bias on the few-shot learning capability. We propose a Prompt-based Decision Transformer (Prompt-DT), which leverages the sequential modeling ability of the Transformer architecture and the prompt framework to achieve few-shot adaptation in offline RL. We design the trajectory prompt, which contains segments of the few-shot demonstrations, and encodes task-specific information to guide policy generation. Our experiments in five MuJoCo control benchmarks show that Prompt-DT is a strong few-shot learner without any extra finetuning on unseen target tasks. Prompt-DT outperforms its variants and strong meta offline RL baselines by a large margin with a trajectory prompt containing only a few timesteps. Prompt-DT is also robust to prompt length changes and can generalize to out-of-distribution (OOD) environments. Project page: \href{https://mxu34.github.io/PromptDT/}{https://mxu34.github.io/PromptDT/}.}
}

Metadata:DOI registry

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Prompting Decision Transformer for Few-Shot Policy Generalization

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Prompt-based Decision Transformer (Prompt-DT) Architecture

Knowl 2 — Trajectory Prompt Representation in Offline RL

Knowl 3 — Prompt-DT Training and Prompt Generation Algorithms

Knowl 4 — Prompt-DT Zero-Gradient Few-Shot Evaluation Algorithm

Knowl 5 — Benchmark Setup and Baseline Configurations for Few-Shot Generalization

Knowl 6 — Few-Shot Generalization Performance Across MuJoCo Benchmarks

Knowl 7 — Effect of Trajectory Prompt Length and Segmentation on Generalization

Knowl 8 — Sensitivity of Prompt-DT to Demonstration and Training Quality

Knowl 9 — Out-of-Distribution Policy Extrapolation in Ant-dir

Knowl 10 — Limitations in Compositional Benchmarks and Prompt Weighting

References

Citation

BibTeX

Source Code

Access the Paper