Prompting Decision Transformer for Few-Shot Policy Generalization

Mengdi Xu; Yikang Shen; Shun Zhang; Yuchen Lu; Ding Zhao; Joshua B. Tenenbaum; Chuang Gan

Source-linked AI summary

Prompting Decision Transformer for Few-Shot Policy Generalization

Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua B. Tenenbaum, Chuang Gan

arXiv:2206.13499v1Jun 27, 2022cs.LGcs.AIcs.CVcs.RO

TL;DR

Offline RL has difficulty generalizing to unseen tasks, while prior meta-RL work emphasizes algorithmic adaptation. Prompt-DT uses Transformer sequence modeling with short trajectory prompts to adapt without finetuning, outperforming strong baselines across five MuJoCo benchmarks and extending to OOD tasks.

Problem
Offline RL struggles to generalize to unseen tasks under distribution shift, motivating investigation of architecture inductive bias for few-shot adaptation.
Method
Prompt-DT trains a Transformer on multi-task offline trajectories and uses a few-shot trajectory prompt to guide policy generation without finetuning.
Results
Across five MuJoCo benchmarks, Prompt-DT outperformed strong meta offline RL baselines without target-task finetuning and generalized to out-of-distribution tasks.
Takeaways & Limitations
Short trajectory prompts can support few-shot policy adaptation in offline RL, with robustness to prompt length but sensitivity to prompt quality.
Takeaways & Limitations
Prompt-DT and Prompt-MT-BC failed to generalize in Meta-World’s ML10 benchmark when prompts were subsampled from expert trajectories.

Abstract

from arXiv · show

Humans 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.

1. Introduction

Offline RL struggles to generalize from training tasks to unseen tasks, motivating Prompt-DT’s architecture-based few-shot adaptation with short trajectory prompts and no target-task finetuning.

Offline RL can fail on unseen tasks because distribution shift makes good test-task policies difficult to find.
Prompt-DT investigates Transformer architecture inductive bias as an alternative to improving few-shot adaptation solely through meta-learning algorithms.
Prompt-DT trains on offline trajectories from multiple tasks and conditions target-trajectory prediction on same-task trajectory prompts.At evaluation, a new task is represented by a handful of trajectories totaling at most 15 steps, with no extra finetuning.
Prompt-DT is proposed as a Transformer-based model that adapts to unseen tasks through short prompts constructed from a handful of trajectories.
Across five MuJoCo benchmarks, Prompt-DT performs strongly without target-task finetuning and beats strong meta offline RL baselines by a large margin.
Prompt-DT is robust to prompt-length changes, sensitive to prompt quality, and generalizes to out-of-distribution tasks where prior methods fail.

2. Related Work

Related work spans offline RL, meta-RL, sequence modeling, few-shot learning, and prompt-based learning as foundations for rapid adaptation from limited data.

Offline Reinforcement Learning: Offline RL learns policies from pre-collected trajectories but must address value estimation and distribution mismatch without online interaction.
Meta-Reinforcement Learning: Meta-RL generalizes knowledge across tasks, commonly using inner-loop adaptation and outer-loop optimization, but MAML-style optimization can be computationally expensive.
Policy Learning as Sequence Modeling: Decision Transformer models offline RL as sequence prediction using state, action, and reward-to-go tokens in a Transformer.
Few-Shot Learning: Few-shot learning aims to generalize rapidly to new tasks from a few supervised samples and can reduce data-gathering effort.
Prompt-based Learning: Prompt-based learning modifies inputs into templates with blanks that language models fill probabilistically to produce desired outputs.

3. Preliminaries

The paper frames RL as sequential decision-making and offline RL as learning from pre-collected data, then represents policies and prompts as Transformer-modeled trajectory sequences.

RL Foundations: An RL task is modeled as an MDP with state and action spaces, transition and reward functions, and an initial-state distribution.
Offline RL: Offline RL learns from pre-collected data because online training can be inefficient, time-consuming, or unsafe in safety-critical environments.
Offline Meta-RL: Offline meta-RL trains on one set of tasks and seeks optimal policies on a disjoint test-task set using task-specific trajectory datasets.
Decision Transformer: Decision Transformer models trajectories with state, action, and reward-to-go tuples, using reward-to-go to guide action selection toward return optimization.
Decision Transformer: A Decision Transformer input contains the most recent K-step history as τ = (r̂_t−K+1, s_t−K+1, a_t−K+1, ..., r̂_t, s_t, a_t).
Prompt-DT: Prompt-DT combines a trajectory prompt with recent history and autoregressively outputs actions corresponding to states in the input sequence.

4. Prompt-based Decision Transformer

Prompt-DT reformulates few-shot policy generalization as prompt-augmented sequence modeling, using short trajectory demonstrations to adapt to unseen tasks without finetuning. Its Transformer processes the prompt with recent history to generate actions.

4. Prompt-based Decision Transformer: Prompt-DT is a Transformer architecture for few-shot policy generalization to unseen tasks using short trajectory prompts.The method is built on Decision Transformer and introduces minimal architectural changes for prompt-augmented sequence modeling.
4.1. Problem Formulation: Offline few-shot RL seeks generalization to new tasks after offline training, without finetuning or gradient updates on task-specific data.The formulation contrasts this goal with offline meta-RL methods that update model weights using task-specific data or interactions.
4.2. Prompt Representation: Trajectory prompts are sequences of state, action, and reward-to-go tuples sampled from a few-shot demonstration dataset to encode task-specific information.They can implicitly capture transition-model and reward-function information while being easier to obtain than text prompts.
4.2. Prompt Representation: K⋆ denotes the number of environment steps in the prompt, which is shorter than the task horizon and identifies the task without providing enough data for imitation.The prompt is therefore intended to specify the task rather than reproduce a complete solution.
4.3. Prompt-DT Architecture: Prompt-DT combines a K⋆-step trajectory prompt with the most recent K-step history and autoregressively predicts actions through a causal Transformer.The resulting input contains 3(K⋆ + K) tokens because each timestep is represented by state, action, and reward-to-go.
4.4. Algorithms: During training, Prompt-DT minimizes mean-squared prediction error for actions in prompt-history pairs sampled across training tasks.Batches aggregate task-level gradient estimates, while stochastic prompts contain J segments of length H with K⋆ = JH.
4.4. Algorithms: Trajectory prompts are generated by sampling episodes and fixed-length segments from task-specific demonstrations, then concatenating those segments.Algorithm 2 uses a task, demonstrations, a sampled episode count, and segmentation length as inputs.
4.4. Algorithms: At evaluation, the pretrained model samples a task prompt, initializes recent history with zeros, and updates that history with streaming environment feedback while generating actions.The evaluation loop uses the prompt and recent context at each timestep without model-weight updates.

5. Experiments

The experiments evaluate few-shot generalization across five MuJoCo control benchmarks using accumulated episode reward and compare Prompt-DT with prompt, multitask, finetuning, and meta-RL baselines. The study examines prompt quantity and quality, task generalization, and OOD transfer.

5. Experiments: The experiments test few-shot policy generalization, prompt quantity and quality, and out-of-distribution task generalization using episode accumulated reward.The evaluation metric is the episode accumulated reward.
5.1. Environments and Datasets: Five MuJoCo control benchmarks comprise Cheetah-dir, Cheetah-vel, Ant-dir, Dial, and Meta-World reach-v2.The environments cover direction, velocity, robotic dialing, and three-dimensional reaching tasks.
5.1. Environments and Datasets: Figure 2 compares Prompt-DT, Prompt-MT-BC, MT-ORL, MT-BC-Finetune, and MACAW on episodic accumulated returns in never-before-seen tasks.All methods use the same expert dataset D, and plots use three seeds.
5.2. Baselines: Prompt-DT removes prompt augmentation from MT-ORL, reward-to-go tokens from Prompt-MT-BC, and both components from MT-BC-Finetune, which adapts through target-task finetuning.These variants isolate the effects of prompts, reward-to-go tokens, and finetuning.
5.2. Baselines: MACAW is the state-of-the-art offline meta-RL comparison method described as sample-efficient and previously superior to multiple finetuning baselines on three control tasks.The comparison includes MACAW alongside three Prompt-DT variants.

6. Discussion

Prompt-DT generalizes from short trajectory prompts, often outperforming adaptation and no-prompt baselines, while prompt quantity and quality shape performance across environments.

Few-shot policy generalization: Prompt-DT and Prompt-MT-BC perform similarly in four environments, indicating trajectory prompts can provide sufficient task-specific information there.In Dial, Prompt-DT outperforms Prompt-MT-BC, suggesting rewards also help when prompts alone are insufficient.
Few-shot policy generalization: Prompt-based methods converge faster than MACAW and, with 5 versus 256 adaptation data points, achieve better asymptotic performance in Cheetah-vel and Ant-dir.They achieve similar performance to MACAW in Cheetah-dir and consistently outperform MT-BC-Finetune.
Prompt quantity: Prompt-DT requires only a small expert prompt, varying the number of episodes and segment length while using a recent history of K = 20 timesteps.The ablation evaluates prompt quantity across three environments with 3 seeds.
Prompt quantity: Nonexpert episodes with negative returns may reduce the effectiveness of prompts constructed from multiple trajectories in Ant-dir.This is presented as a conjecture about the observed prompt-quantity effect.
Prompt quantity: Prompt-DT is generally insensitive to prompt quantity and can extract task-specific information from prompts containing only a few timesteps.The experiments suggest sequential prompt information is not crucial in the tested MuJoCo control settings.
Prompt quality: Prompt-DT adjusts generated actions to prompt quality: random prompts lower returns, while expert or medium prompts yield higher performance when training data are expert or medium quality.The ablation uses expert, medium, and random datasets and demonstrations in Cheetah-vel.
Out-of-distribution generalization: Prompt-DT performs better than no-prompt baselines on Ant-dir tasks whose goal directions lie outside the training task range.The evaluation uses eight training tasks and three testing tasks, with two below and one above the training index range; large variance is reported.

7. Conclusion

The paper proposes Prompt-DT for offline few-shot reinforcement learning and reports strong benchmark performance with prompt-based adaptation. It presents sequence-prediction modeling as a simple alternative to separately trained actor-critic networks, while identifying failures on complex compositional tasks.

Conclusion: Prompt-DT is proposed to solve offline few-shot reinforcement learning and outperforms MACAW across multiple benchmark domains.The conclusion reports empirical evaluation against the state-of-the-art offline meta-RL algorithm MACAW.
Conclusion: Prompt-DT is robust to prompt-length changes when trained with expert data but remains sensitive to the quality of prompt data.This conclusion summarizes the prompt quantity and prompt quality ablations.
Conclusion: The work presents sequence-prediction models as a first application in offline few-shot reinforcement learning and uses a prompt-based Transformer instead of separate policy and value networks.The authors describe the algorithm as simple to implement relative to MACAW’s actor-critic training.
Limitations and future work: Prompt-DT and Prompt-MT-BC fail to generalize in Meta-World’s ML10 benchmark when prompts are subsampled from expert trajectories.The authors identify better prompts and prompt-based algorithms for complex compositional tasks as future work.

A. Hyperparameters

The appendix reports shared hyperparameters for Prompt-DT, Prompt-MT-BC, MT-ORL, and MT-BC-Finetune, alongside environment-specific settings for the prompt-based methods and MACAW settings.

Table 2 lists common hyperparameters shared by Prompt-DT, Prompt-MT-BC, MT-ORL, and MT-BC-Finetune.
Table 3 lists environment-specific hyperparameters for Prompt-DT and Prompt-MT-BC.
Table 4 lists MACAW’s hyperparameters and their values.

B. Experiment Details

The experiments document separate training/testing task indexes for evaluating in-distribution and out-of-distribution generalization.

Table 5 specifies training and testing task indexes for in-distribution generalization experiments.
Table 6 specifies training and testing task indexes for out-of-distribution generalization experiments.

C. Ablation Study

The prompt-length ablation evaluates Prompt-DT and Prompt-MT-BC across different prompt quantities and lengths, finding limited sensitivity in three environments.

Prompt-length ablations evaluate Prompt-DT and Prompt-MT-BC using training curves across varying prompt quantities.
In Cheetah-dir, Cheetah-vel, and Ant-dir, prompt-based methods are not sensitive to the number of episodes or prompt length.
Figure 5 reports the effect of trajectory prompt length on Prompt-DT’s performance, using 3 seeds per plot.
Figure 6 reports the effect of trajectory prompt length on Prompt-MT-BC’s performance, using 3 seeds per plot.

C.2. The Effect of Finetune Data’s Quantity on MT-BC-Finetune

The ablation examines how finetuning-data quantity affects MT-BC-Finetune and contrasts limited-data instability with adaptation using more data and steps.

With limited finetune data, MT-BC-Finetune struggles to adapt to every task.
MT-BC-Finetune shows large performance variance in Cheetah-dir with adaptation batch size 256, but smaller variances in Cheetah-vel and Ant-dir.
The Cheetah-dir variance may result from disjoint state distributions across tasks with opposite rewards by design.
With 100 finetune steps and adaptation batch size 1280, MT-BC-Finetune can adapt to test tasks.
Figure 7 varies MT-BC-Finetune’s finetune-data quantity using batch sizes 256 and 128.
Figure 8 reports the effect of finetune-data quantity on MACAW, with MACAW-256 and MACAW-128 configurations.

C.3. The Effect of Finetune Data’s Quantity on MACAW

Increasing MACAW’s adaptation batch size consistently improves performance, but MACAW-256 still underperforms Prompt-DT in several environments and converges more slowly in Meta-World reach-v2.

C.3. The Effect of Finetune Data’s Quantity on MACAW: Increasing MACAW’s adaptation batch size consistently improves performance across environments.Each curve uses 10 finetuning gradient steps.
C.3. The Effect of Finetune Data’s Quantity on MACAW: MACAW-256 still underperforms Prompt-DT in Cheetah-dir, Cheetah-vel, and Ant-dir.
C.3. The Effect of Finetune Data’s Quantity on MACAW: In Meta-World reach-v2, MACAW-256 reaches similar asymptotic performance to Prompt-DT but converges more slowly.