Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement
Authors: Zhi Wang
Li Zhang
Wenhao Wu
Yuanheng Zhu
Dongbin Zhao
Chunlin Chen
Venue: NeurIPS 2024 poster (2024)
Keywords: decision transformer
offline meta reinforcement learning
world model
Primary area: reinforcement_learning
TL;DR: We leverage the sequential modeling ability of the transformer architecture and robust task representation learning via world model disentanglement to achieve efficient generalization in offline meta-RL.

Abstract: A longstanding goal of artificial general intelligence is highly capable generalists that can learn from diverse experiences and generalize to unseen tasks. The language and vision communities have seen remarkable progress toward this trend by scaling up transformer-based models trained on massive datasets, while reinforcement learning (RL) agents still suffer from poor generalization capacity under such paradigms. To tackle this challenge, we propose Meta Decision Transformer (Meta-DT), which leverages the sequential modeling ability of the transformer architecture and robust task representation learning via world model disentanglement to achieve efficient generalization in offline meta-RL. We pretrain a context-aware world model to learn a compact task representation, and inject it as a contextual condition to the causal transformer to guide task-oriented sequence generation. Then, we subtly utilize history trajectories generated by the meta-policy as a self-guided prompt to exploit the architectural inductive bias. We select the trajectory segment that yields the largest prediction error on the pretrained world model to construct the prompt, aiming to encode task-specific information complementary to the world model maximally. Notably, the proposed framework eliminates the requirement of any expert demonstration or domain knowledge at test time. Experimental results on MuJoCo and Meta-World benchmarks across various dataset types show that Meta-DT exhibits superior few and zero-shot generalization capacity compared to strong baselines while being more practical with fewer prerequisites. Our code is available at https://github.com/NJU-RL/Meta-DT.

Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement
Zhi Wang'
Li Zhang' Wenhao Wu' Yuanheng Zhu? Dongbin Zhao? Chunlin Chen'
'Nanjing University Institutiio Automation, Chinese Academy of Sciences edu.
fzhiwang, clchen) edu. en (yuanheng. zhu, dongbin 2hao]@ia..... on
{lizhang, whao_ (amooww nju. edu. cn
Abstract
A longstanding goal of artificial general intelligence is highly capable generalists that can learn from diverse experiences and generalize to unseen tasks. The lan- guage and vision communities have seen remarkable progress toward this trend by scaling up transformer- based models trained on massive datasets, while rein- forcement learning (RL) agents still suffer from poor generalization capacity under such paradigms. To tackle this challenge, we propose Meta Decision Transformer Meta--D), which leverages the sequential modeling ability of the transformer chitecture and robust task representation learning via world model disentanglement achieve efficient generalization in offline meta-RI.. We pretrain context-aware world model to learn a compact task representation, and inject it as contextual con- dition the causal transformer to guide task-oriented sequence generation. Then, we subtly utilize history trajectories generated by the meta-policy as a self-guided prompt to exploit the architectural inductive bias. We select the trajectory segment that yields the largest prediction eгrоr on the pretrained world model to construct the prompt, aiming to encode task speciseei information complementary to the world model maximally. Notably, the proposed framework eliminates the requirement of any expert demonstration or domain knowledge test time. Experimental results on MuloCo and Werld benchmarks across various dataset types show that Meta-DTT exhibits superior few and zero-shot generalization capacity compared strong baselines while being more practical with fewer prerequisites. Our code is available httptp///gtthubbb..//// wttpoeuuwooooo//-R//u------e
Introduction
Building generalist models that solve various tasks by training on vast task-agnostic datasets has emerged as dominant paradigm in the language [I] and vision [ communities. Offline rein- forcement learning (RL) ( allows learning an optimal policy from trajectories collected by some behavior policies without access to the environments, which holds tremendous promise for turning massive datasets into powerful generic decision-making engines [] 6). akin the rise of large language models (LLMs) like GPT 7 A performant example is the decision transformer (DT) g that leverages the transformer's strong sequence modeling capability trajectory data and achieves promising results on offline RL [9 12].
Transformer- baased based large models have shown remarkable scaling properties and high transferability across various domains, including language modeling [13] 77, computer vision [5), and image generation [6 17). However, their counterparts RL are usually specialized a narrowly defined task and struggle with poor generalization to unseen tasks, due to distribution shift and the lack of
CCorrespondence Chunlin Chen <elchen @ njj.eed.cnn.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

self-supervised pretraining techniques [8, 19 20). To promote generalization, Prompt-DT 21 uses expert demonstrations as a high-quality prompt to encode task-specific information. Generalized DT [22] conditions on some hindsight information, e.g., statistics of the reward distribution, to match task-specific feature distributions. These works usually rely on domain knowledge at test time, e.g.. expert demonstrations or hindsight statistics, which is expensive and even infeasible to acquire in advance for unseen tasks [23). The aforementioned limitations raise a key question: Can we design an offline meta-RL framework to achieve efficient generalization to unseen tasks while drawing upon advances in the sequence modeling paradigm with the scalable transformer architecture?
In offline RL, the collected dataset depends on the task and behavior policy. When behavior policies are highly correlated with tasks in the training dataset, the agent tends memorize the features of behavior policies and produce biased task inference at test time due to the change of behavior policies (4 25]. One major challenge for generalization is how to accurately encode task-relevant information for extrapolating knowledge across tasks, while minimizing the requirement on the distribution of pre-collected data and behavior policies. RL agents typically learn through active interaction with the environment, and the transition dynamics p(r, s' s a) completely describes the characteristics the underlying environment 226.. The environment dynamics, also called world model 2 28 is intrinsically invariant to behavior policies or collected datasets, thus presenting promising alternative for robustly encoding task beliefs. By capturing compact representations of the environment, world models carry the potential for substantial transfer between tasks (29] continual learning [30]. and generalization from offline datasets [3
Inspired by this, we propose novel framework for offline meta-RI., named Meta-DTT that leverages robust task representation learning via world model disentanglement to conduct task-oriented se- quence generation for efficient generalization. First, we pretrain a context-aware world model to capture task-relevant information from the multi-task offline dataset. The world model contains an en- coder that abstracts dynamics-specific information into a compact task representation, and a decoder that predicts the reward and state transition functions conditioned on that representation. Second, we inject the task representation as contextual label to the transformer to guide task-conditioned sequence generation. In this way. the autoregressive model learns to estimate the conditional output of multi-task distributions and achieve desired returns based on the task- context. Finally, we leverage the past trajectories generated by the meta- peea-p as a self-guided prompt to exploit the architecture inductive bias, akin to Prompt-DTT (21). We feed available trajectory segments to the pretrained world model and choose the segment with the largest prediction error to construct the prompt, aiming encode task task-specific information complementary to the world model maximally.
In summary, our main contributions are as follows:
Generalization Ability. We propose Meta-DT to leverage the conditional sequence modeling paradigm with the scalable transformer architecture to achieve efficient generalization across unseen tasks without any expert demonstration or domain knowledge at test time.
Context-Aware World Model. We introduce a context-aware world model to learn a compact task representation capable of generalizing across a distribution of varying environment dynamics. Complementary Prompt. We design a self-guided prompt that encodes task-specific information complementary to the world model maximally, harnessing the architectural inductive bias.
Superior Performance. Experiments on various benchmarks show that Meta-DT exhibits higher few and zero-shot generalization capacity, while being more practical with fewer prerequisites.
2 Related Work
Offline Meta-RL. Offline RL 333 allows learning optimal policies from pre-collected offline datasets without online interactions with the environment [34 35). Offline meta-RL learns to generalize to new tasks via training on a distribution of such offline tasks (25 36.. Optimization-
meta-learning seek policy parameter initialization that requires only a few adaptation steps to new tasks. MACAW 39 introduces this architecture into based RL subroutines, and uses simple supervised regression objectives for sample-efficient offline meta- veeta-training. On the other hand, context-based meta-learning methods learn a context encoder to perform approximate inference on task representations, and condition the meta-policy on the approximate belief for generalization, such as PEARL. 40, VariBAD 41, FOCAL [2), CORRO (24, CSRO 25), and UNICORN [4 These works extended from the online meta-RI RL setting are mainly trained by temporal difference
methods [7 38

---

learning, the dominant paradigm in RL 19 44]. This paradigm might be prone to instabilities due to function approximation, off-policy learning, and bootstrapping. together known as the deadly triad [26]. Moreover, many of these works rely on hand-designed heuristics to keep the policy within the offline dataset distribution 45). This motivates us to turn to harness the lens of conditional sequence modeling with the transformer architecture to scale existing offline meta-RI. algorithms.
RL as Sequence Modeling, The recent rapid progress in (self) supervised learning models is in large part predicted by empirical scaling laws with the transformer architecture, such as GPT 77, ViT [14, and DiT [46.. Decision transformer [8] first introduces the transformer's 's sequence modeling capacity to solving offline RL without temporal difference learning, which autoregressively outputs optimal actions conditioned on desired returns. As concurrent study, trajectory transformer [47 directly models distributions over sequences of states, actions, and rewards, followed by beam search as planning algorithm a model-based manner. Extending the LL---like structure RI., this paradigm activates a new pathway toward scaling powerful RI. engines with large-scale compute and data [48]
For multi-task learning, Gato 49 trains multi-modal generalist policy in the GPT style with decoder-only transformer. Multi-game DT 4 trains a suite of up to 46 Atari games simultancously, and controls policy generation inference time with a pretrained classifier that indicates the expert level of behaviors. For generalization to unseen tasks, Prompt-DT (21] exploits prefix prompt architecture and Generalized DT [22] designs hindsight information matching structure to encoder task-specific information. These works usually require domain knowledge at test time, such as expert demonstrations or hindsight statistics. TºGDT [50] conditions action generation on three-tier guidance of global transitions, local relevant experiences, and vectorized time embedding. CMT [18] provides a pretrain and prompt-tuning paradigm where a task prompt is extracted from offline trajectories with an encoder-only transformer to realize few-shot generalization.
3 Preliminaries
3.1 Offline Meta Reinforcement Learning
RL is generally formulated as Markov decision process (MDP) M= A,T, where S/A the statc/action space, T' R is the state transition/reward functiion and y is the discount. The objective is to find policy T(ass) to maximize the expected return as max,, JM - EE R(s, at)].
In offline meta-RL, we assume that the task follows distribution M, = (S. A, R. P(M). where tasks share the same state-action spaces while vary in the reward and state transition functions, i.e. enviromment dynamics. For each task from total of N training tasks { [ an offline dataset D;= [$ ot.j. J )/ is collected by an arbitrary behavior policy z/. The agent can only access the offline datasets {Dr} to train a meta-policy Tmeta At test time with the few-shot setting, given an unseen task MM- ~ P(M). the agent can access a small dataset D = {(3j. Tj. s s')}? to construct the task prompt or infer the task belief before policy evaluation. For the zero-shot setting, the trained meta- meta-policy is directly evaluated on the unseen task by interacting with the test environment to estimate the expected return. The objective is learn meta-policy to maximize the expected episodic return over test tasks as J m = EMP)(() [Jar (Tmeaa)].
3.2 Decision Transformer
By leveraging the superiority of the attention mechanism [13] in extracting dependencies between sequences, transformers have gained substantial popularity in the language and vision communities [7] 14. Decision transformer
[ introduces associated advances to solve RL, problems and re- re-frames sequential modeling problem, inheriting the transformer's s efficiency and scalability when modeling long trajectory data. DT models the probability of the next sequence token He conditioned on all tokens prior to it:
offline RL as a cettur-conditions
similar to contemporary decoder-only sequence models [51]. The sequence we consider has the form: z = (... Rí, St, a,, ".. ). where R is the agent's 's target return for the rest of the trajectory. Such a sequence order respects the causal structure of the decision process. When training with offline data, RG = > Ti, and during testing, R R==G* Ei=o Fi, where where G is the target return for an entire episode. The same timestep embedding is concatenated to the embeddings of Ar 8t, and de, and the head corresponding to the state token is trained to predict the action by minimizing its error from the ground truth.
3

---

Multi-Task Offline Dataset p ~ Dj
K-K+1
H
Meta Decision Transformer
1.1.
Context-Aware World Model
Context Encoder (x((((
6K
R
Complementary Prompt
Reward Decoder R St, ati z()
Transition Decoder
Tp($1+118t agi zt)
jtk MAX Σ
History Trajectories
[(v,, o, sp)
Figure I: The overview of Meta-DT.. We pretrain context-aware world model to accurately disentan- gle task-specific information. contains a context encoder Ep that abstracts recent h-step history μ into a compact task representation z/, and the generalized decoders (Rg. T.p) that predict the reward and next state conditioned on zt. Then, the inferred task representation injected as contextual condition to the causal transformer to guide task-orienter sequence generation. Finally, we design self-guided prompt from history trajectories generated by the meta-policy test time. We select the trajectory segment that yields the largest prediction error on the pretrained world model, aiming to encode fask-relevant information complementary to the world model maximally.
4 Method
In this section, we present Meta Decision Transformer (Meta-DT)), a novel offline meta-RI. framework that draws upon advances in conditional sequence modeling with robust task representation learning and a complementary prompt design for efficient generalization across unseen tasks. Fig. B iillustaates the method overview, and the following subsections present the detailed implementations.
4.1 Context-Aware World Model
The key to efficient generalization is how to accurately encode task-relevant information for extrapo- lating knowledge across tasks, which proves learner is passively presented with a fixed dataset and aims to recognize the pattern within that dataset. In contrast, RL typically learns from active interactions with the environment, and aims at optimal policy that receives the maximal return in the environment, rather than only recognizing the pattern within the pre-collected dataset. Since the offline dataset depends on both the task and the behavior policy, the task information could be entangled with the features of behavior policies, thus producing biased task inference at test time due to the shift in behavior policies. Taking 2D navigation as an example, tasks differ in goals and the behavior policy is going towards the goal for each task. The algorithm can easily distinguish tasks based on state-action distributions rather than reward functions, leading to extrapolating errors when the behavior policy shifts to random exploration during testing.
challenging in RL. A supervised unsupervised
To address this intrinsic challenge, ought to disentangle task-relevant information from behav- ior policies. The transition dynamics, i.e., the reward and state transition functions p(s', rrss, a), completely describes the characteristics of the underlying environment. Naturally, this environ- ment dynamics, also called world model, keeps invariant to behavior policies or collected datasets, and could be promising alternative for accurately inferring task beliefs. Therefore, introduce a context-aware world model to learn robust task representations that can generalize to unseen environments with varying transition dynamics.
In the typical meta-learning setting, the reward and state transition functions that are unique to cach MDP are unknown, but also share some common structure across the task distribution P(M). There exists a true variable that represents either a task description or task identity, but we do not have access to this information. Hence, a given task M,, the reward and state transition functions can be approximated by generalized context-aware world model W that shared across tasks as
latent representation
approximate that variable. For
Wi(re, S141 ax) ~ (17,,,4411,,
(I)

---

Since we do not have access to the true task description or identity, we need to infer task representation z) given the agent's h-step experience up to timestep collected in task M, as
p (8t-h: h,, Gt-h: Ft- 8t-1, (1-1, (0t=1,7t-1,,81
(2)
The intuition is that the true task belief of the underlying MDP can be abstracted from the agent's recent experiences, analogous to recent studies in meta-RL 41 36).
We separate the reasoning about the world model into two parts: i) encoding the synamicc--ppetiee information into latent task representation, and ii) decoding the environment dynamics conditioned on that representation. First, we use a simple yet effective context encoder Eg to embed recent experiences into task representation as 2 = Egg. (xz) Second, the decoder contains a generalized reward model Rg and state transition model Tw The task representation is augmented into the input of both models to predict the instant reward fz+1 = (s;, ay; 2)) and next state $ =T (St, ag; z(). Under the assumption that tasks with similar contexts will behave similarily [52 53], the proposed world model can extrapolate the meta-level knowledge across tasks by accurately capturing task- relevant information from training environments. The context encoder is jointly trained by minimizing the reward and state transition prediction error conditioned on the task representation as
[ 0 ()) = Eug~D. [Be [re41 - Ra ae: 2)))² + (sL+1 - TTp ar z)) "] (3)
4.2
Complementary Prompt
Recent works suggest the prompt framework as an effective paradigm for pretraining transformer- based large models on massive datasets and adapting them to new scenarios with few or no labeled data [54). Prompt-DT (2I) adopts that paradigm to RL problems by prepending a task prompt to the DT's input. At test time, the agent assumed to access handful of expert demonstrations to construct the prompt and perform few-shot generalization to new tasks. However, in weal-world scenarios, it is expensive and even infeasible to acquire such domain knowledge as expert demonstrations in advance for unseen tasks. Especially in RL, agents typically learn from interacting with initially unknown environment. Recent studies 33 23 also suggest that the quality of demonstrations must be high enough to act as well-constructed prompt. Otherwise, the performance of Prompt-DT may degrade since some medium or random data can easily disrupt the abstraction of task-specific information.
Here, we design self-guided prompt to achieve more realistic generalization at test time. In the ideal case where the testing policy is optimal. past experiences during evaluation can act as high-qualiny demonstrations to construct the prompt. Motivated by this, we leverage history trajectories generated by the meta-policy during evaluation as a self-guided prompt to enjoy the power of architecture inductive bias, while eliminating the dependency on expensive domain knowledge. Though, the meta-policy may generate medium or even inferior data during initial training. As meta-learning proceeds, the policy can exhibit increasing generalization capacity via the context-aware world model, thus generating trajectories that gradually approach expert demonstrations.
Both the world model (algorithmic perspective) and the self-guided prompt (architecture perspec- tive) aim to extract task specific information to guide policy generation across tasks. To facilitate collaboration between these two parts, we force the prompt to maximally complement the world model towards the same goal. We feed all available segments selected from history trajectories to the pretrained world model, and use the segment with the largest prediction error to construct the prompt. In this way, the complementary prompt attempts to encode the portion of task-relevant information that the world model struggles to capture effectively.
Formally, the complementary prompt is sequence segment containing multiple tuples, (R¹, s" aa), which are sampled from the agent's history trajectories. Each element with superscript * associated with the prompt. For given task M,, we obtain a k-step prompt T as
T/' = (Rz, sj aj. R Rj $/+1, aj aj.,...,, Rj S74R dj+R):
1+h و
[ Ro axs = (8e+1 Tan--Ten TTpl-, (1s =f)
where
max
Following Prompt-DT, we choose a much smaller value of k than the task horizon, ensuring that the prompt only contains the information needed to help identify the task but insufficient information for the agent to imitate. Compared to the world model that explicitly learns the reward and transition functions, the complementary prompt store partial information to implicitly capture task dynamics.
5

---

4.3 Meta-DT Architecture
Our Meta-DT architecture is built on decision transformer g and solves the offline meta-RL problem through the lens of conditional sequence modeling with robust task representation learning. We first pretrain the context-aware world model and keep fixed to produce compact task representations for the downstream meta-training process. For each task M,, we use the context encoder from the pretrained world model to infer the contextual information for each timestep as z Fuy (14). Then, the input of Meta-DT DT consists of two parts: i) the k-step complementary prompt 7; derived from and ii) the most recent K step history T augmented by the learned task representations as
Tt - Re- k+. 81- K+1, at- K+.. 24, St, St, 0z)
(5)
The input sequence (T (T)',.") corresponds to 3k + 4K tokens in the transformer, and Meta-DT autoregressively outputs k+ + K actions at heads corresponding to state tokens in the input sequence. During training, we construct the prompt from the top few trajectories that obtain the highest returns in the offline dataset. Meta-DT minimizes errors between the predicted and real actions in the data for the K stetep history. At test time with a few- setting, the meta meta-policy is allowed to interact with the environment for few episodes, and Meta-DT utilizes these history interactions to construct the self-guided prompt. For the zero-shot setting, we ablate the prompt component and directly evaluate Meta-DDT on test tasks. Corresponding algorithm pseudocodes are given in Appendix A.
Experiments
We comprehensively evaluate the generalization capacity of Meta-DT on popular benchmarking domains across various dataset types. In general, we aim to answer the following questions:
Can Meta-DT achieve consistent performance gain on the few and zero-shot generalization to unseen tasks compared with other strong baselines? (Secs. 5. and 5.2)
How do the context-aware world model, the self-guided prompt design, and the complementary prompt construction affect the generalization performance, respectively? (Sec. 5.3)
Is Mcta-DT robust to the quality of offline datasets? (Sec. 5.4
Environments. We evaluate all tested methods on three classical benchmarks in meta-RL: i) the 2D navigation environment Point-Robot 255 ii) the multi-task MuloCo control [55] 36). containing Cheetah-Vel, Cheetah Ant-Dir, Hopper-Param, and Walker-Param; and iii) the Meta- World manipulation platform [56]. including Reach, Sweep, and Door-Lock. For each environment, we randomly sample distribution of tasks and divide them into the training t and test set Ties. On each training task, we use SAC (57) to train a single-task policy independently for collecting datasets. We consider three ways construct offline datasets: Medium, Mixed, and Expert. More details about environments and datasets are given in Appendix B.
Baselines. We compare Meta-D DT to four competitive baselines that cover two distinctive paradigms in offline meta-RL: the based 1) Prompt-Di [21]. 2) Generalized DT [22], and the temporal
difference-based 3) CORRO 24 4) CSRO 25 More details about baselines are given in Appendix C.
All evaluated methods are carried out with different random seeds, and the mean of the received return is plotted with 95% bootstrapped confidence intervals of the mean (shaded). The standard errors are presented for numerical results. Appendix D gives implementation details of Meta-DT. Appendix E presents hyperparameter analysis on the context horizon h, the prompt length k and the number of training tasks. Appendix F shows experimental results on Meta- World domains.
5.1 Few-shot Generalization Performance
Note that these baselines work under the few-shot setting, since they require expert demonstrations as task prompts or some wtrr-start data infer the task representation. We first compare Meta- DT to them under an aligned few-shot setting, where each method can leverage a fixed number of trajectories for task inference. Fig. 2 and Table present the testing curves and converged performance of Meta-DT and baselines on various domains using Medium datasets under aligned few-shot setting. these environments with varying reward functions or state transition dynamies, Meta-DT consistently obtains superior performance regarding data efficiency and final asymptotic

---

CORRD
Prompt-D1
Generalized DT
CSRD
Meta-DT (Ours)
Cheetah
Cheetah
-10
-100
600
Returns
-15
Returs
-150
Returns
400
-20
-200
200
-25
0
8
10
-250
0
8
10
0
0
8
10
Timesteps (1o4)
Timesteps (1o4)
Timesteps (1o4)
Ant-Dir
Hogper-Param
Walker-Param
450
400
350
400
Returns
300
Returns
300
Retums
350
300
250
200
250
10
40
50
200
8
10
150
8
10
Timesteps (1e4)
Timesteps (1e4)
Timesteps (1e4)
Figure The received return curves averaged test tasks of Meta-DT and baselines using Medium datasets under an aligned few-shot setting.
Table
Few-shot test returns of Meta-DT DT against baselines using Medium datasets.
Environment
Prompt-DT
Generalized DT
CORRO
CSRO
Meta-DT
Point-Robot
-12.58
-1..97
1.65
-10.18 -1693
Cheetah- Vel
1135.89=1999
-123.47=6.88
-154.12428.83
11.1.83+9.16
-99.28
Cheetah-Dir
5..22
272.94-12.002
566.08
556.97114.59
608.17 14.18
Ant-Dir
287.23+
268.53
449.3
412.00 231.66+30.01
Hopper-Param
309.90.5.74
12.87
332.37+12.90
348.20 332.84
Walker-Param
18.33
368.07
368.02
388.08
11.11
results. A noteworthy point is that Meta-DT outperforms baselines by a larger margin in relatively complex environments like Ant-Dir, which demonstrates the high generalization capacity of Meta-DT when tackling challenging problems. Moreover, Meta-DT generally exhibits a lower variance during learning, signifying not only superior learning efficiency but also enhanced training stability.
5.2 Zero-shot Generalization Performance
Since Meta-DT can derive real-time task representations from its htstep experience μ4 via the pre- trained context encoder, can achieve zero-shot policy adaptation to unseen tasks. To demonstrate its zero-shot generalization ability, we ablate the prompt component and directly evaluate Meta-DT on test tasks. For a fair comparison, we modify the baselines to an aligned zero-shot setting, where prompt or warm-start data is inaccessible before policy evaluation. All methods can only use samples generated by the trained meta-policy during evaluation to construct task prompts (Prompt-DT), calculate hindsight statistics (Generalized DT),, or infer task representations (CORRO and CSRO).
Fig. 3 and pressen the testing curves and converged performance of Meta-DT and baselines on various domains using Medium datasets under an aligned zero-shot setting. Unsurprisingly, most baselines exhibit collapsed performance (about 20%-35% drop on average) compared their few- shot counterparts. The main reason is that they usually require expert demonstrations as task prompts or some high-quality warmup data accurately capture task information. contrast, Meta-DT can abstract compact task presentations via the context-aware world model, and its performance in zero-shot scenarios drops merely little (about 5% on average) compared to few-shot counterpart. The superiority of Meta-DT over all baselines is more pronounced when deployed to zero-shot adaptation scenarios. In addition to its enhanced generalization capacity, this result also demonstrates the superior practicability of Meta-D DT with fewer prerequisites in real-world applications.
7

---

CORRD
Prompt-DT
Generalized DT
CSRD
Meta-DT (Ours)
Cheetah
Cheetah
600
-100
-15
400
Returns
-20
Returss
-150
Returns
200
-25
200
-30
-250
0
-39
0
8
10
0
8
10
0
8
10
Timesteps (1o4)
Timesteps (1o4)
Timesteps (1o4)
Ant-Dir
Hogper-Param
Walker-Param
450
400
350
400
Regums
Returss
300
Retums
350
300
200
250
250
100
0
10
40
50
200
10
150
8
10
Timesteps (1e4)
Timesteps (1e4)
Timesteps (1e4)
Figure 3: The received return curves averaged over test tasks of Meta-DT and baselines using Medium datasets under an aligned zero-shot setting.
Table 2: Zero-shot test returns of Meta-DT against baselines using Medium datasets. The 4 denotes the performance drop compared to the few-shot setting.
Environment
Prompt-DT
Generalized DT
CORRO
CSRO
Meta-DT
Robot
-17.3 (138%)
10. (113%)
1.6.3
-23.0 (140%)
-10.7 +0. (15%)
Cheetah- Vel
-148.3 (19%)
(19%)
-16..149..4488))
-140.4 (126%)
-100.1 =28 (119)
Cheetah-Do
212.4=:185.2 (164%)
129.2=106:5
(177%)
146.4 (174%)
200.9
542.4 (L11%)
Ant-Dir
150.6=:5:4444885 ((48%)
162.2 (140%)
235.5 (1.06
(16%)
Happer-Param 301.7:1122 (L38)
318.7
320.0.1.0.0 (14%)
329.9
169 (L19)
342.0
142.0=174(2255
Walker-Param 317.9 117.9 (L119))
355.2 14.0 14:0
342.4131.1 (17%)
347.9 17799122.1 (110%)
397.4 (125)
5.3
Ablation Study
To test the respective contributions of each component, compare Meta-DT four ablations: w/o //_ context, only removes the task representation 2 from the input of the casual transformer; ii) w/o_com, oomom it randomly chooses segment from candidate trajectory to construct the prompt, i.e. removing the complementary way for constructing the prompt; iii) //o__prompt it directly removes the prompt component; iv) DT, it removes all components and degrades the original DT. For ablations, all other structural modules are kept consistent strictly with the full Meta-DT.
Fig. 4 and Table show ablation results representative domains using Medium datasets. First, Meta-DT obtains a decreased performance when any component is removed, which demonstrates that all components
essential for Meta-DT's capability and they complement each other. Second, ablating the context incurs a more significant performance drop than ablating the complementary way or the whole prompt. It indicates that task representation learning via the world model plays more vital role capturing task-relevant information. Third, the performance order woo_rropp < w/o_com fo_ Meta-DT shows that using the self-guided prompt can improve the performance, and leveraging the complementary way to construct the prompt can further achieve a performance gain. Another interesting point is that the performance gain achieved by the self-guided prompt and the complementary way more significant in complex environments like Ant-Dir than in simpler ones like Point-Robot. again verifies our superiority tackling challenging tasks.
5.4 Robustness the Quality of Offline Datasets
A good offline algorithm ought to be robust to different types of datasets that involve wide range of behavioral policies (58] To test how Meta-DTT performs with data of different qualities, we also conduct experiments on the Mixed and Expert datasets. Figs. present test return curves of
8

---

DT
w/ context
W/O prompt
wo complementary
Meto-DT
Point-Robot
Cheetah-DDr
-10
600
-12
Returss
Retums
3000
-14
200
-16
200
100
%
0.2
0.4
0.8
1.0
Timesteps (1e4)
Timesteps (1e5)
Timesteps (1e5)
Figure 4: Test return curves of Meta-DT ablations using Medium datasets. w/o_contex removes task representation, w/o_com removes the complementary way, and w/o__pompp removes the prompt.
Table 3: Test returns of Meta-DT ablations using Medium datasets.
Environment
w/o_com
w/o_ w/o_prompt
DT
Meta-DT
Point-Robot
-16.44 1.93
05.5
-10.61 +0.15
-1..04
0..1
Cheetah-Dir
551.30
580.34+33..8
542.44 13.84
008.18+4.1
Ant-Dir
54.36
390.70
368.94 10.56
412.00 15.364
CORRO
Prompt-DT
Generalized
Meta-DT (Ours)
Point-Robot
Cheetah-Vel
Cheetah-Dir
-75
600
-100
Retumss
Retumss
Returss
125
400
150
200
-175
10
30
Timesteps (1e4)
Timesteps (1e4)
Timesteps (1e4)
Ant: Dir
Hopper-Param
Walker Param
1000
375
500
350
450
600
Returns
325
Returns
400
350
400
300
200
275
250
200
Timesteps (1e4)
Timesteps (1e4)
Timesteps (1e4)
Figure 5: Few-shot test curves of Meta-DT and baselines using Mixed datasets.
Meta-DT and baselines with few-shot setting, and Table 4 shows corresponding converged results. In Mixed datasets, Meta-DT still outperforms all baselines by a large margin, especially complex environments like Ant-Dir and Walker-Param.
In Expert datasets, Meta-DT exhibits superior generalization capacity than the three baselines of Generalized DT, CORRO, and CSRO. Compared to Prompt-DT, Meta-DT obtains significantly better performance in Robnt-Robot, Cheetah-Vel. and Ant-Dir, and obtains comparable performance in the other three environments. This phenomenon is because Prompt-DT is sensitive to the quality of prompt demonstrations, The performance can drop when provided with prompts from Medium or Mixed datasets, also mentioned in the original paper [2]] and subsequent studies (36). Hence, Prompt-DT can achieve satisfactory performance only when expert demonstrations are available at test time. When tackling offline tasks with lower data qualities, Prompt-DT would induce significant performance gap compared to Meta-DT. In summary, the above results verify that our method robust the dataset quality and is more practical with fewer prerequisites in real-world scenarios.

---

CORRO
Prompt-OT
Generalized DT
CSRO
Meta-DT (Ours)
Point-Robot
Cheetah-Vel
-5
-50
1000
Returns
Returns
-100
Returns
800
-15
600
-150
400
-200
200
-25
10
20
30
6
0
10
10
20
20
Timesteps (1e4)
Timesteps (1e4)
Timesteps (1e4)
Hopper-Param
Walker-Param
1000
400
450
800
400
Returss
Retums
350
Retumns
330
300
300
250
200
10
20
30
40
50
250
6
10
200
8
10
Timesteps (1e4)
Timesteps (1e4)
Timesteps (1e4)
Figure 6: Few-shot test curves of Meta-DT and baselines using Expert datasets.
Table 4: Few-shot test returns of Meta-DT against baselines using Mixed and Expert datasets.
Mixed
Prompt-DT
Generalized DT
CORRO
CSRO
Meta-DT
Point-Robot
+.55
01551
110.88+1.22
3.771
-8.39 +.28
Cheetah- Vel
19..34=14.87
-86.52=10.62
-81.59437.17
-1.0.55
-79.90
Cheetah-Dir
376.66
480.49+48.38
27..80+24.87
631.34 579.90+
Ant-Dir
269.58=22.64
511.97
255.49132.64
330.15133.77
908.48 20.73
Hopper-Paran Param
320.76
338.36 14.07
335.12.426.73
351.96
358.05
10.69
Walker-Param
20.22
394.32 19.78
330.54114.07
334.34122.57
470.36 12.85
Expert
Prompt-DT
Generalized DT
CORRO
CSRO
Meta-DT
Point-Robot
+0.4
-13.99
-7.76=0.18
-19.4222.10
-6.1
Cheetah- Vel
-133.78 18.24
-62.95
-111.47
-129.00-22.24
-52.42
Cheetah-Dir
960.32+
806.30
628.64461.11
621.05
47..41
Ant-Dir
678.07 68.74
149.52
18..42-13.83
417.37
961.27
Hopper-Param
39..79+11.44
358.56+11.75
338.17+47.36
383.51 358.29
Walker-Param
447.1
421.96:14.770
352.02
336.89+16.71
137.79+18.21
6 Conclusions, Limitations, and Future Work
In the paper, we tackle the offline meta-RL challenge via leveraging advances in sequence modeling with scalable transformers, marking step toward developing highly capable generalists akin to those in the language and vision communities. Improvements in few and zero-shot generalization capacities highlight the potential impact of our robust task representation learning and self-guided complementary prompt design. Without requiring any expert demonstration or domain knowledge at test time, our method exhibits superior practicability with fewer prerequisites in real-world scenarios.
Though, our method requires a two-phase process pre- pre-training the world model and training the causal transformer. A promising improvement is to develop unified framework that simultancously abstracts the task representation and learns the meta- veta--poicy, akin to in-context learning [59]. Also, generalist model is trained relatively lightweight datasets compared to popular large models. crucial future step is to deploy our model on significantly larger datasets, unlocking the scaling law with the transformer architecture. This also aligns with the urgent trend that RL practitioners striving to break through. Another interesting line to leverage self-supervised learning [5 60 to facilitate task representation learning at scale. We leave these future work.
10

---

Acknowledgements
We thank Zican Hu and Zhenhong Sun for their helpful discussions. This work was supported by the National Natural Science Foundation of China (Nos. 62376122 & 6207316) the Nanjing University Integrated Research Platform of the Ministry of Education -Top Talents Program, the Beijing Natural Science Foundation under Grant 4232056, the Beijing Nova Program under Grant 20240484514 14, and the Youth Innovation Promotion Association CAS under Grant 2021132.
References
[I] Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. BERT:
Pre-training
of deep bidirectional transformers for language understanding. In Proceedings of the North American Chapter of the Association for Computational Linguistics, pages 4171-4186, 2019.
[2] Anurag Arab,, Mostafa Dehghami, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. ViVIT: A video vision transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6836-6846, 2021.
[3] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proceedings of International Conference on Machine Learning, pages 2052-2062, 2019.
[4] Kuang- Huei Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. In Advances in Neural Information Processing Systems, volume 35, pages 27921-27936, 2022.
[5] Thomas Schmied, Markus Hofmarcher, Fabian Paischer, Razvan Pascanu, and Sepp Hochreiter. Learning to modulate pre-trained models in rl. Advances in Neural Information Processing Systems, 2023.
[6] Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Proceedings of Conference on Robot Learming, pages 785-799, 2023,
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901, 2020.
[8] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, volume 34, pages 15084-15097, 2021.
[9] David Brandfonbrennrr, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does retum--onditionee supervised learning work for offline reinforcement learning? In Advances in Neural Information Processing Systems, volume 35, pages 1542-1553, 2022.
[10] Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. In Advances in Neural Information Processing Systems, 2023.
[11] Anirudhan Badrinath, Yannis Flet-Berliac, Allen Nie, and Emma Brunskill. Waypoint trans- former: Reinforcement learning via supervised learning with intermediate targets. In Advances in Neural Information Processing Systems, volume 36, 2023.
[12] Prajjwal Bhargava, Rohan Chitnis, Alborz Geramifard, Shagun Sodhani, and Amy Zhang. When should we prefer decision transformers for offline reinforcement learning? In Proceedings of International Conference on Learning Representations, 2024.
[13] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Informa- tion Processing Systems, pages 5998-6008, 2017.

---

[14] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. In Proceedings of International Conference on Learning Representations, 2021.
[15] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16000-16009, 2022.
[16] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Iya Sutskever. Zero-shot text-to-image generation. In Proceedings of International Conference on Machine Learning, pages 8821-8831, 2021.
[17] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, et al. Scaling autoregressive models for content-rich text-to-image generation. Transactions on Machine Learning Research, 2022.
[18] Runji Lin, Ye Li, Xidong Feng. Zhaowei Zhang, Xian Hong Wu Fung, Haifeng Zhang, Jun Wang, Yali Du, and Yaodong Yang. Contextual transformer for offiin meta reinforcement learning. In NeurlPS 2022 Foundation Models for Decision Making Workshop, 2022.
[19] Shengehao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt-tuning decision transformer with preference ranking. arXiv preprint arXiv:2305.09644 2023.
[20] Zhihui Xie, Zichuan Lin, Deheng Ye, Qiang Fu, Yang Wei, and Shuai Li. Future-conditioned unsupervised pretraining for decision transformer. In Proceedings of International Conference on Machine Learning, pages 38187-38203 2023.
[21] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In Proceedings of International Conference Machine Learning, volume 162, pages 24631-24645, 2022.
[22] Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. In Proceedings of International Conference on Learning Representations, 2022.
[23] Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang. Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. In Advances in Neural Information Processing Systems, 2023.
[24] Haoqi Yuan and Zongqing Lu. Robust task representations for offline menf-reiinornemment learning via contrastive learning. In Proceedings of International Conference on Machine Learning, pages 25747-25759, 2022.
[25] Yunkai Gao, Rui Zhang, Jiaming Guo, Fan Wu, Qi Yi, Shaohui Peng, Siming Lan, Ruizhi Chen, Zidong Du, Xing Hu, et al. Context shift reduction for offline reitforeiffrr learning. In Advances in Neural Information Processing Systems, volume 36, 2023.
[26] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018.
[27] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122 2018.
[28] Tim Brooks, Bill Peebles, Connor Homes, Will DePue, Yufei Guo, et al. Video generation models as world simulators. 2024.
[29] Arunkumar Byravan, Jost Tobias Springenberg, Abbas Abdolmalcki, Roland Hafner, Michael Neunert, Thomas Lampe, Noah Siegel, Nicolas Heess, and Martin Riedmiller. Imagined value gradients: Model-based policy optimization with tranferable latent dynamics models. In Proceedings of the Conference on Robot Learning, volume 100, pages 566-589, 2020.
[30] Samuel Kessler, Mateusz Ostaszewski, Michat PaweiBortkiewwczz Mateusz Žarski, Maciej Wolczyk, Jack Parker-Holder Stephen J. Roberts, and Piotr Miš. The effectiveness of world models for continual reinforcement learning. Proceedings of Conference on Lifelong Learning Agents, volume 232, pages 184-204, 2023.
12

---

[31] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in Neural Information Processing Systems, volume 34, pages 18954-28998 2021.
[32] Zohar Rimon, Tom Jurgenson, Orr Krupnik, Gilad Adler, and Aviv Tamar. MAMBA: an effective world model approach for meta-reinforcemmnt learning. Proceedings of International Conference Leurning Representations, 2024.
[33] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint wrXiv:2055.0664 2020.
[34] Alexander Nikulin, Vladislav Kurenkov, Denis Tarasov, and Sergey Kolesnikov. Anti- exploration by random network distillation. In Proceedings of International Conference on Machine Learning, 2023.
[35] Yuhang Ran, Yi-Chen Li, Fuxiang Zhang, Zongzhang Zhang, and Yang Yu. Policy regularization with dataset constraint for offline reinforcement learning. In Proceedings of International Conference on Machine Learning, 2023.
[36] Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang, MetaDif- fuser: Diffusion model as conditional planner for offline meta-rl. In Proceedings of International Conference on Machine Learning. 2023.
[37] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adapta- tion of deep networks. Proceedings of International conference on machine learning, pages 1 1126-1135. 2017.
[38] Zhongwen Xu, Hado P van Hasselt, and David Silver. Meta-gradien gradient reinforcement learning. In Advances in Neural Information Processing Systems, volume 31, 2018.
[39] Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In Proceedings of International Conference on Machine Learning, pages 7780-7790 2021.
[40] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcemmnt learning via probabílistic context variables. In Proceedings of International Conference on Machine Learning, pages 5331-5340, 2019.
[41] Luisa Zintgraf. Sebastian Schulze, Cong Lu, Leo Feng, Maximilian Igl, Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson. VariBAD: Variational Bayes-adaptive deep RL via meta-learning. The Journal of Machine Learning Research, 22(1):1119- 2021.
[42] Lanqing Li, Rui Yang, and Dijun Luo. FOCAL: Efficient fully-offline meta-reinforcemmnn learning via distance metric learning and behavior regularization. Proceedings of International Conference Learning Representations, 2021.
[43] Lanqing Li, Hai Zhang, Xinyu Zhang, Shatong Zhu, Junqiao Zhao, and Pheng-Am Heng. Towards an information theoretic framework of context- based offline meta-reiiforremenn learning, arXiv preprint arXiv:240..02429, 2024.
[44] Zhi Wang, Chunlin Chen, and Daoyi Dong. Lifelong incremental reinforcement learning with online Bayesian inference. IEEE Transactions on Neural Networks and Learning Systems, -308)-400344016 2022.
[45] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? In Proceedings of International Conference on Learning Representations, 2023.
[46] William Peebles and Saining Xie. Scalable diffusion models with transformers. Proceedings of the JEEE/CVF International Conference on Computer Vision, pages 4195-4203 2023.
[47] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems, volume 34, pages 1273-1286 2021.
13

---

[48] Linghui Meng, Muning Wen, Chenyang Le, Xiyun Li, Dengpeng Xing, Weinan Zhang, Ying Wen, Haifeng Zhang, Jun Wang. Yaodong Yang, and Bo Xu. Offline pre-trained multi-agent decision transformer. Machine Intelligence Research, 20(2):333-445 2023.
[49] Scott Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel martt, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg. Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. Transactions on Machine Learning Research, 2022.
[50] Zhe Wang, Haozhu Wang, and Yanjun (Jane) Qi. T3GDT: Three-tier tokens guide decision transformer for offline meta reinforcement learning. In NeurlPS 2023 Workshop on Robot Learning, 2023.
[51] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, et al. PaLM: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
[52] Aditya Modi, Nan Jiang, Satinder Singh, and Ambuj Tewari. Markov decision processes with continuous side information. In Proceedings of Algorithmic Learning Theory, pages 597-618, 2018.
[53] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In Proceedings of International Conference on Machine Learning, pages 5757-5766, 2020.
[54] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 55(9))1-35, 2023.
[55] Emanuel Todorov, Tom Erez, and Yuval Tassa. MuJoCo: A physics engine for model-based control. In Proceedings of IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033, 2012.
[56] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta- Meta-World: benchmark and evaluation for multi-task and meta reinforcement learning. In Proceedings of Conference Robot Learning, pages 1094- 1 100, 2020.
[57] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of International Conference on Machine Learning, pages 1861-1870, 2018.
[58] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072119, 2020.
[59] Jake Grigsby, Linxi Fan, and Yuke Zhu. AMAGO: Scalable in-context reinforcement learning for adaptive agents. In Proceedings of International Conference on Learning Representations, 2024.
[60] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018
14

---

Appendix A. Algorithm Pesudocodes
Based on the implementations in Sec. 4, this section gives the brief procedures our method. First, Algorithm Baathmmmppeeeeett the pretraining of the context-aware world model. Then, Algorithm shows the pipeline of training Meta-DT, where the sub- uub-procedure generating the complementary prompt is given in Algorithm [] Finally, Algorithm and Algorithm 5 show the few-shot and zero-shot evaluations on test tasks, respectively.
Algorithm 1: Pretraining the Context-Awart World Model
Input:
Training tasks T and corresponding offline datasets Dirain Context encoder Ey: Reward decoder Rs; Transition decoder T Batch size B
Context horizon h;
for each iteration do
for bo 1,.., B do
Sample task M, ~ Tirain and obtain the corresponding dataset Dj from Dirain Sample whole trajectory (sh: ab. rô, sf, af ri.... from D;
Sample transition tuple (st, (e, Tt; 86+1) with randomly selected
Obtain its h-step history μu} = (St-hy at- Tt-hs 8t-1, at 1,-1, 8t)
Compute the context z = (1()
Compute the predicted reward fr = (sr, ari 2/) and next state $ = Top (sts ari
end
10
Update Ev, Rs, and Txp using the loss as
C(V, 25 pp) Σ [ro+1 Rg (8t, 41z =1)) + (se41 Tople, u4i =1))
end
Algorithm 2: Training Meta Decision Transformer
Input:
Training tasks Tiaain and corresponding offline datasets Dirain Causal transformer Fe; Pretrained context encoder Es and decoders R2, T.p Context horizon h;
Trajectory length K; Batch size B
for each task M, Train do Construct the demo dataset Delemio using the top few trajectories with the highest returns in Dj
3 end
for each iteration do for b 1,2,.... do
Sample task M, Train and obtain the corresponding dataset D; from Dirain
Sample trajectory of length K from D,, T; {R, a) K+1
Infer the context of each timestep & from its h-step history using the pretrained context encoder as 2; Eg (St-k, at-h: Ft-h: .... $t-1. Qt 0
Augment the trajectory T; with per-step context as T,+ = {(2, R, (,s,,) a))o-)]K- K+
Sample prompt GetPrompt (Ddlmmo, Eg, Rs Tx) as in Algorithm Get input r inpur 7
10
11
end
Get batch B [T, input B pred
13
Jbml imput
Fe (T: imput) VT;
B
15
(()) Σ - Snme
2
16
Few or Zero-Shot Evaluation along with training
17 end
15

---

Algorithm 3: Complementary Prompt Generation (GetPrompt)
Input:
Demo Dataset De Prompt length A:
Pretrained context encoder Exp and decoders Rg. T
Sample trajectory (sp, dg, ro,..., ..., St, at, rt; .t.) from Delmm
Get the segment with the largest prediction error on the pretrained world model as
max, Ei (1e+1 (5t, azj =)))² + (st+1 - Top (st. azi =f)) where af = B (T()
Obtain the k-step prompt T. (2j 87, 03, Ri+1 01+1 Ri +Rx a
Algorithm 4: Meta-DT Few-Shot Evaluation on Test Tasks
Input: Test tasks T
Target return G";
Evaluation episodes N Prompt length
Context horizon h;
Trajectory length K;
Trained meta-DT Fo:
Pretrained context encoder Ey and decoders Rg. Txp
for each task M, T do Initialize the demo dataset D.lemo for n 1,.... N do
Initialize desired return = G for each timestep do
Infer the context from h-step history using Eg as
2 = Eng (8t- H1+1, - #+1,.. 81 at 1 11- 8c) Augment the trajectory т; with context as T. = {(z, R, a) K+1
Sample prompt T. Prompo (Doomo, Exp, R2, Tx) as in Algorithm
Get action Fo((TT"
Step env. and get feedback 8, a, r, R + R -
10
11
Append (R, a) recent history т
12
end
Append the whole trajectory т; to the demo dataset as Delemo Delemo U
13
14
end
15 end
Algorithm 5: Meta-DT Zhro-Shot Evaluation on Test Tasks
Input: Test tasks T
Target return G*
Context horizon h;
Trajectory length K
Trained meta-DT Fo: Pretrained context encoder Exp
for each task M4 T do Initialize desired return A G
for each timestep t do
Infer the context from h-step history using Exp
2 =
(st- 1+1, Gt- H41, Tt- #+1,... 8t-1, at-1, 11-1,
Augment the trajectory TÜ with context 7, (, Get action (T(TT)
Step env. and get feedback s, a, r, R + R R-
Append (R, s, a) to recent history т
end
10 end
16

---

Appendix B. The Details of Environments and Dataset Construction
In this section, we show details of evaluation environments over a variety of testbeds, as well as the offline dataset collection process conducted on these environments.
B.1. The Details of Environments
Following the mainstream studies in offline meta-RL. 33
24
we adopt three classical benchmarks: and the Meta-War Geta-World [56]. We
the 2D navigation 25), the multi-task MuloCo control 555 evaluate all tested methods on the following environments
36
Point-Robot: a problem of point agent navigating to given goal position in the 2D space. The observation is the 2D coordinate of the robot and the goal is not included in the observation. The action space 0. 0.1]² with each dimension corresponding to the moving distance in the horizontal and vertical directions. The reward function is defined as the negative distance between the point agent and the goal location. Each learning episode always starts from the fixed origin and terminates at the horizon of 20, Tasks differ in goal positions that are uniformly distributed in a unit square, resulting in the variation of the reward functions.
Cheetah- Vel. Cheetah-Dir. Ant-Dir: multi-task MuJoCo continuous control benchmarks in which the reward functions differ across tasks. Cheetah- Vel requires a planar cheetah robot to run at a particular velocity in the positive x-direction, and the reward function is negatively correlated with the absolute value between the current velocity of the agent and goal. Cheetah-Dir//Aee -ir is to control a 2D cheetah/3D quadruped ant robot to move in a specific direction, and the reward function is the cosine product the agent's 's velocity and the goal direction. The goal velocity and direction are uniformly sampled from the distribution 00.07,,, 3.0] for Cheetah-Ves and U0, 2T] for Ant-Dir. In Cheetah-Dir the goal directions are limited forward and backward. The maximal episode step is set to 200.
Hopper-Param, Walker-Paramt multi-task MuJoCo benchmarks where tasks differ in state transition dynamics. The two benchmarks control one-legged hopper and a two-legged walker robot to run as fast as possible. The reward function is proportional to the rumning velocity in the positive z-direction, which remains consistent for different tasks. In both domains, the physical parameters of body mass, inertia, damping, and friction are randomized across tasks. The agent needs to move forward with varying environment dynamics to accomplish the task. The maximal episode step is set to 200 for both.
Reach, Sweep, Door-Lock: three typical environments from the robotic manipulation benchmark Woold. Reach, Sweep, and Door-Lock control a robotic arm to reach a goal location in 3D space, to sweep a puck off the table, and to lock the door by rotating the lock clockwise, respectively. Tasks are assigned with random goal object positions for Reach, with random puck positions for Sweep, and with random door positions for Door-Lock. The variations in goal/puck/door positions are not provided in the observation, forcing meta-RL. RII algorithms to adapt to the goal through trial-and-error. The maximal episode step is set to 500.
For the Point-Robot and MuJoCo environments, sample 45 tasks for training and another 5 held-out tasks for testing. For Meta- World environments, we sample 15 tasks for training and held-out tasks for testing.
B.2. The Details of Datasets
We employ the soft actor-critic (SAC) [7 algorithm to train policy independently for each task. Table5]shows the detailed hyperparameters for training the SAC agents each environment. Figs. 7 show the learning curves of independently training the SAC agents for each training and testing task in all evaluation domains. During training, we periodically save the policy checkpoints to generate various types of offline datasets as
Medium: using medium policy that achieves one-third one-second score of an expert policy.
All trajectories are generated by this medium policy.
using expert policy that obtains the optimal return in the environment. In practice, we load the last policy checkpoint after the training has converged, and use this expert policy to generate all trajectories.
Expert:
17

---

Mixed: mixed dataset of the Medium and Expert types. We use a combination of the medium (70%) and expert (30%) datasets collected as above generate the mixed dataset.
Specifically, we generate 100 trajectories for the Point-Robo and MuJoCo environments, and 300 trajectories for the Meta-World environments, using the saved checkpoint policies to construct corresponding offline datasets.
B.3. Further discussions
The above environments are relatively homogenous in nature. Our generalist model is trained on relatively lightweight datasets compared to popular large models. The urgent trend that RL practitioners are striving to break through is to deploy significantly large datasets with more diversity, unlocking the scaling law with the transformer architecture. It would be interesting to see how Meta-DT generalizes across worlds and tasks with more diversity, like training on K-levels of an Atari game and generalizing to the remaining N K levels of well-suited Atari games.
The challenges in extending to such paradigms might include: i) how to collect large-scale datasets with sufficient diversity and high quality, how tackle the heterogenous task diversity in hard cases, and iii) how to scale RL models to very large network architectures like in large language or visual models. To tackle these potential challenges, we conjecture several promising solutions including: i) leveraging self-suuervised learning to facilitate task representation learning at scale, ii) enabling efficient prompt design or prompt tuning to facilitate in-context RL, and iii) incorporating effective network architectures like mixture-of-experts to handle heterogenous domain diversity.
Table 5: Hyperparameters of SAC used to collect multi-task datasets.
Training
Warmup steps
Save frequency
Learning rate
Soft update
Discount factor
Entropy ratio
Environments
steps
Point-Robot
2000
100
40
3e-4
0.005
0.99
0.2
Cheetah-Dii
500000
10000
10000
3e-4
0.005
0.99
0.2
Cheetah- Vel
500000
2000
10000
3e-4
0.005
0.99
0.2
Ant-Dir
500000
2000
10000
3e-4
0,005
0.99
0.2
Hopper-Param
500000
10000
10000
3e-4
0.005
0.99
0.2
Walker-Param
500000
2000
10000
3e-4
0.005
0.99
0.2
Reach 200000
5000
1000
3e-4
0.005
0.99
0.2
Sweep
200000
5000
1000
3e-4
0.005
0.99
0.2
Door-Lock
200000
5000
1000
3e-4
0.005
0.99
0.2
Appendix C. The Details of Baselines
This section gives the details of the representative five baselines, including three DT-hased approaches, one temporal difference-batte approach, and one diffusion- basod approach. These baselines are thoughtfully selected to encompass the three distinctive categories offline meta-RL methods. Also, since our method Meta-DT belongs the DT-based category, we adopt more approaches from this kind as our baselines. The baselines are introduced as follows:
Prompt-DT (11] is DTT-ased method that leverages the sequential modeling ability of the Transformer architecture and the prompt framework to achieve few-shot adaptation in offline RL. It designs the trajectory prompt, which contains segments of the few-shot demonstrations, and encodes task-specific information to guide policy generation. At test time, it assumes that the agent assess handful of expert demonstrations to construct the prompt.
Generalized DT (22, is a DT-based method that formulates multi-task learning as hindsight information matching (HIM) problem, i.e., training policies that can output the rest of the trajectory to match some statistics of future information. It defines offline multi-task state- marginal matching and imitation learning as two generic HIM problems to evaluate the proposed Categorical DT and Bi-directional DT.
18

---

Point Robat (Train Tasks)
Half-Chestah-Vd Vel (Train Tasks)
Half Chlftch Oir (Train Tasks)
-100
1250
Returs
750
-20
4400
500
-25
5500
250
Timesteps
Clost Half-Cheetah-vell [Teet Tasles)
Cheesahhtt Taskes)
1250
1000
Returrns
-200
Recurns
750
-15
500
-300
250
-400
0.0
0.5
1.0
2.0
Timesteps (1e3)
Timesteps (1e5)
Timesteps (1e5)
Ant-Dir (Tiain Tasks)
CeopperRanc-Paraannee [Tain Tasks)
Moliker-Rand-Paraane (Taiin Tasks)
1200
1000
000
1000
Returns
Returss
600
Returss
800
400
600
250
200
-250
0
Timesteps (1eb)
Timesteps
(1e5)
Ant-Dir (lest Tasks)
Hopper- Ranee--Ran--ae Params (lesst Teoks)
Welker-Rand Params (lest Teoks]
1000
1000
600
800
800
Raturns
Returss
Returns
600
a00
800
400
300
400
200
200
100
Timesteps (3e5)
Timestops (1e5)
Timeszops (1051
Figure 7: Learning curves of training independent SAC agents in Point-Robot and MuJoCo domains.
Reach
Sweep
Tasks)
5000
4000
3000
3000
Returs
3000
Returss
Returns
2000
2000
2000
1000
1000
0.0
0.0
0.5
1.0
1.5
2.0
Timesteas (1e5)
Timesteps
(1e5)
Timesteps (1e5)
(Test
DoorLock
5000
1000
4000
4000
3000
3000
Returss
Returns
2000
2000
1000
0.0
0.0
0.3
1.0
2.0
Timestegs (1e5)
Timesteps (3e5)
Timesteps (1e5)
Figure 8: Learning curves of independently training the SAC agents for each task in Meta-World.
19

---

CORRO [24]. is a temporal difference- method with a contrastive learning framework to learn task representations that are robust to the distribution mismatch of behavior policies in training and testing. It formalizes the learning objective mutual information maximization between the representation and task, to maximally eliminate the influence of behavior policies from task representations. At test time, it needs to collect context trajectory with an arbitrary policy to infer the task representation.
CSRO [25] is temporal difference- based method that proposes a max- max-min mutual information representation learning framework to mitigate the distribution discrepancy between the contexts used for training (from the behavior policy) and testing (for the exploration policy). At test time, it explores the new task with a non-prior strategy and collects few-shot context data to infer the task representation. It demonstrates superior performance than prior context-based meta-RI. methods.
Note that these baselines can only work under the few-shot setting. since they require several expert demonstrations as the task prompt or some warm-start data infer the task representation. We first compare Meta-DT to them under an aligned few-shot setting, where each method can leverage fixed number of trajectories for task inference.
Further, we demonstrate the zero-shot generalization capability of Meta-DT and baselines by directly evaluating the meta-policy on the unseen task without collecting any data in advance. For Meta-D'T, we ablate the prompt component and derive real-time task representations from the history trajectory via the pre-trained context encoder. For fair comparison, we modify the baselines an aligned zero-shot setting, where prompt demonstrations or warm-start data are inaccessible before policy evaluation. All these baselines can only use samples generated by the trained meta-policy during evaluation to construct task prompts (Prompt-DT), calculate hindsight statistics (Generalized DT), or infer task representations (CORRO and CSRO).
Appendix D. Implementation Details of Meta-DT
Network Architecture. In this paper, we use simple network structures for the context-aware world model, including the context encoder, the reward decoder, and the state transition decoder. Specifically, the context encoder contains a fully-connected multi-layer perceptron (MLP) and GRU network with ReLU the activation function. The GRU cell encodes the agent's h-step history (st-n, at-h, r ..... 8t-1, 4, Pt-1, 8t) into a 128-dimensional vector, and the MLP transforms this vector into 16-dimensional embedding, i.e., the task representation The reward decoder is an MLP that takes as input the (s, a, s', 2) tuple and predicts the reward r through two 128ddimensional hidden layers. Similarly, the state transition decoder is an MLP that takes input the (s,a,r, 2) tuple and predicts the next state s' through two 128-dimensional hidden layers.
We implement the proposed Meta-DT based on the official codebase released by DT S htttps//giiuub..om///////
We follow most of the hyperparameters as they did. Specifically, task representations 2, returns R, states 8, and actions a are fed into modality-speciff lin- embeddings, and the same positional episodic timestep encoding is added tokens corresponding to the same timestep. Tokens are fed into a GPT architecture which predicts actions autoregressively using a causal self- attention mask. In summary, Table shows the details of network structures.
Algorithm Hyperparameters. We evaluate the proposed Meta-DT algorithm on various testbeds with different types of offline datasets. Each report unit involves running one environment (Point- Robot, Cheetah- Vel. Cheetah-Din Dir, Ant-Dir, Hopper- Hopper-Param, Walker-Parre (Medium, Expert, Mixed), Some common hy hyperparameters across all report units are set optimizer Adam, weight decay Ic-4, linear warmup steps for learning rate decay 10000, gradient norm clip 0.25, dropout 0, and batch size 128. Table presents the detailed hyperparameters of Meta-DT trained on the Point-Robe and MuJoCo domains with the Medium, Expert, and Mixed datasets. Table X presents the detailed hyperparameters of Meta-DT -T trained on Meta- World environments with the Medium datasets.
Reach) with one dataset
Compute. We train our models one Nvidia RTX4080 GPU with the Intel Core 19-10900X CPU and 256G RAM. Pretraining the context-aware world model and the causal transformer costs about 0.5-6 hours and 1-8 hours, respectively, depending on the complexity of the environment.
20

---

Table 6: The network configurations used for Meta-DT.
World Model
Value
Causal Transformer
Value
GRU hidden dim
128
layers num
3
task representation dim decoder hidden dim
16
attention heads num
1
128
embedding dim
128
decoder hidden layers num activation function
2
activation function
ReLU
ReLU
Table Hyperparameters of Meta-DT on Point-Robot and MuJoCo domains with various datasets.
Medium
Point-Robo Cheetah- Vel Cheetah-Dir Ant-Dir Hopper-Param Walker-Parne
training steps
100000
100000
100000
500000
100000
100000
sequence length K
20
20
20
20
context horizon
4
learning rate
le-4
5e-5
6c-6
5c-5
5e-5
5e-5
prompt length Ic
5
5
5
5
Expert
Point-Robe Cheetah- Vel Cheetah- Dir Ant-Dir Hopper-
Walker-Param
training steps
300000
200000
300000
500000
100000
100000
sequence length K
20
20
20
20
20
context horizon h
4
4
4
4
4
4
learning rate
le-4
5e-5
5e-5
5e-5
5e-5
5e-5
prompt length &
5
5
Mixed training steps
Robtt Cheetah- Vel Cheetah- Dir Ant-Dir Hopper-
Hopper-Param Param Walker-Parne
300000
200000
300000
500000
100000
100000
sequence length K 8
20
20
20
20
20
context horizon h
4
4
4
4
4
learning rate
le-4
5e-5
5e-5
5e-5
5e-5
5e-5
prompt length
5
5
5
Table 8: Hyperparameters of Meta-DT trained on Meta-World domains with Medium datasets.
Hyperparameters
Reach
Sweep
Door-Lock
training steps
500000
500000
500000
sequence length K
20
20
20
context horizon h
4
4
learning rate
1e-7
1e-6
le-6
prompt length k
5
Appendix E. Hyperparameter Analysis
Context horizon h. As shown by the ablation studies in Sec. 5.3 the world model disentanglement plays a more vital role in capturing task-relevant information in Meta-DT. The task representation 2 is abstracted from the agent's 's h-step experience μ = (St-h, Q .... St-1, 44-1, 8t)1
task representation learning via
8t). The context horizon h is a key hyperparameter in learning effective task representations. Hence, we conduct additional experiments to analyze the influence of context horizon / Meta-DT's performance.
Fig. 9 and Table present the testing curves and converged performance of Meta-DT with different values of context horizon h on representative environments. Generally, the performance of Mc
21

---

is not sensitive to the pre-defined value of the context horizon. In simple environments like Point- Robot, increasing the context horizon can degrade the final performance a little bit. The most likely reason for this phenomenon is that the context-aware world model might be overfitted with large value of context horizon. In more complex environments like Cheetah-Dir and Ant-Dir, adopting different values of context horizon leads to very close performance. Hence, throughout the paper, we choose the context horizon as 4 for our method to exploit lightweight network design while not sacrificing performance.
horizon-4
horizon-6
horizon-8
Point: Robot
Cheetah Dir
Ant Dir
-10
600
400
Retumss
Retums
500
Retums
350
400
300
-14
300
200
250
-16
1.0
100
Timesteps (1e5)
Timesteps (1e5)
Timesteps (1e5)
Figure 9: The received return curves averaged over test tasks of Meta-DT with different values of context horizon h using Medium datasets under an aligned few-shot setting.
Table 9:
Few-shot test returns averaged over test tasks of Meta-DT with different values of context horizon h using Medium datasets.
Context Horizon
4
Point-Robot
10.18 0.18
-11.08 0.11
-11.81 0.31
Cheetah-Dir
608.18 4.18
606.81 11.34
610.63=6.13
Ant-Dir
412.00 11.53
406.01 = 12.61
401.36-6.65
Prompt length kt. Analogously, we also investigate the influence of prompt length on Meta-DT's performance. Fig. [0] and Table T0 present the testing curves and converged performance of Meta-DDT with different values of prompt length on representative environments. Generally, the performance of Meta-DT is not sensitive the pre-defined value of prompt length.
prompt-3
prompt-5
prompt-7
700
450
-10
600
400
Returns
Returns
-12
350
400
-14
300
200
250
-16
0.2
0.4
1.0
100
0.0
0.2
0.4
0.8
1.0
200
Timesteps (1e5)
Timesteps (1e5)
Timesteps (1e5)
Figure 10: Few-shot return curves of Meta-DT with different prompt length k using Medium datasets.
Table 10: Few-shot test returns of Meta-DT with different prompt length & using Medium datasets.
prompt length &
3
5
7
Point-Robot
-1..18
11..19
-10.22
Cheetah-Dir
606.6912.99
008.1814.11
605.09 1.311
Ant-Dir
408.97+18.01
412.00
413.09
22

---

The number of training tasks. We evaluate Meta-DD DT with a different number of training tasks. Fig. IT and Table [T] present the testing curves and converged performance of Meta-DT with different numbers of training tasks representative environments using Medium datasets. Generally, increasing the number of training tasks can usually improve the generalization ability to test tasks, especially in harder tasks like Ant-Dir.
The result also matches the intuition of function approximation in machine learning. The "gener- alization" refers to the question: How can experience with limited subset of the data space be usefully generalized produce good approximation over a much larger subset? the single-task setting, increasing the valid sample points (before being overfitted) can generally improve the func- tion approximation of the sample space and the generalization to unseen samples at testing. In the meta-learnn veea-learning setting, increasing the valid task points (before being overfitted) can usually boost the function approximation of the task space and the generalization to unseen tasks testing.
n. n_task-15
n. n_task-30
n. task-45
Point: Point-Robot
Cheetah-Vel
Ant-Dir
450
-10
-100
400
Returns
Returns
Returns
.55
-14
-200
250
-16
0.2
0.4
1.0
-250
0.2
0.4
1.0
200
Timesteps (1e5)
Timesteps (1e5)
Timestops (1e5)
Figure 11: Few-shot return curves of Meta-DT with different numbers of training tasks (Medium).
Table 11: Few-shot test returns of Meta-DT with different numbers of training tasks (Medium).
Number of training tasks
15
30
45
Point-Robot
-..220
-1..45
-10.18 0.188
Cheetah- Vel
-85.42
.95.1411..8
-99.28 3.39
Ant-Dir
388.50=9.02
00.071=5.00
Appendix F. Experimental Results on Meta-Worke
Here, present experimental results the robotic manipulation benchmark Meta- WorldWorldd [56]. including three environments of Reach, Sweep, and Door-Lock. Fig. 12 and Table 12 show the testing curves and converged performance of Meta-DT and baselines using Medium datasets under an aligned few-shot setting. The results show that Meta-DT still performs better than the baselines in these challenging and realistic Meta-World environments.
CORRO
Prompt-DT
Generalized DT
CSRO
Meta-DT (Ours)
Reach
Sweep
DoorLock
3500
800
2500
600
2500
Returns
2000
Returns
Returns
AM
2000
400
1500
200
1000
1000
10000.0
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
500
500
0.2
0.4
0.6
0.8
Timestops (1e5)
Timesteps (1e5)
Timesteps (1e5)
Figure 12: Few-shot return curves of MT and baselines Meta-World using Medium datasets.
23

---

Table 12: Few-shot test returns of Meta-DT and baselines on Meta-World using Medium datasets.
Env. Prompt-DT
Generalized DT
CORRO
CSRO
Meta-DT
Reach
2296.55166...
2241.00:1:94.37
1496.941402.21
2298.69
2458.83 162.10
Sweep
549.17 192.10
272.67 177.41
446.50 191.46 760.18 1.887
Door-Lock 2284.95
26.1.24144442
1615.66:124.00 2216.67 116.84 3025.92 = 23.41
Appendix G. Generalization to Out-of-Distribution Tasks
Most meta-RL. studies follow the experimental setting as training in a large distribution of tasks and evaluating few held-out tasks. The held-out tasks have goals (target position or direction) within the goal range calibrated by training tasks. We desire to test whether Meta-DT enables knowledge extrapolation when handling tasks with goals out of the training range, i.e., the generalization ability in out-of-distribution (OOD) tasks. Hence, we conduct experiments generalization ability to OOD tasks. In this case, we sample training tasks with a goal direction of (10, 1.5x] and then test on tasks of [1.5m, 2T) Fig. a3d and Table 13 show the few-shot return curves and converged performance on OOD test tasks using Medium datasets from Ant-Dir.
Ant-Dir to evaluate Meta-DDT's
Obviously, Meta-DT can still obtain better performance than baselines on OOD test tasks, which again verifies our superiority. Our method extrapolates the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model intrinsically invariant to behavior policies collected datasets. For tasks like Ant-Dir, the world model shares some common structure across the task distribution (even for OOD tasks), e.g., the kinematics principle or locomotion skills. Hence, the extrapolation of the world model also works for OOD test tasks in this case.
Ant-Dir
CERRO
400
allaed
Table 13:
Few-shot returns of OOD test tasks using Medium datasets from Ant-Dir.
Retums
300
Method
OOD Performance
200
Prompt-DT
281.73
100
Generalized DT
148,83
CORRO
141.791194.0
0
CSRO
208.92+
Timesteps (1e5)
Meta-DT
412.67 5:515
Figure 13: Few-shot return curves of OOD test tasks using Medium datasets from Ant-Dir.
24

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Please refer to the Abstract and Section Introduction.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Please refer Section 6. Conclusions, Limitations, and Future Work.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: Please refer to Appendix A for our algorithm pseudocodes, Appendix B for the dataset construction, and Appendix for the implementation details of our method. We have also provided our source code at github for reliable reproduction.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We have provided the source code with sufficient instructions at github:
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please refer to Section 5: Experiments, Appendix B, Appendix D, and the source code at github.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Please refer Section Experiments.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please refer Appendix D.
25

---

9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:ptaw..accap......pppoaaaaaa
Question:
Answer: [Yes]
Justification: Please refer to our submitted source code at github.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Please refer to Section 6: Conclusions, Limitations, and Future Work.
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
26


Reviews:
[{"review_id": "MxbrytVQo8", "reviewer": "Reviewer_iZYy", "summary": "This paper proposes a novel Meta-DT method that leverages the task representation from the world model disentanglement. Compared with the previous works, the expert demonstration is not necessary. This method could get the task representation from the trained encoder which is used as the guidance for the autoregressive training. This paper also incorporates an additional prompt as the complementary. The empirical study shows the good performance of the Meta-DT compared with baselines.", "questions": "- Can you explain more about the training hours and parameters for this method?\n- Can you explain the experiments why the method does not outperform the baselines in the expert demonstration?", "limitations": "See weaknesses and questions.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper proposes a novel Meta-DT that disentangles the task information from the trajectories as the prompt in Meta RL tasks.\n- This paper could outperform the baselines in most of the settings, especially on low-quality demonstration sets.", "weaknesses": "- This paper needs to train a world model which is used to represent the dynamic information for the task. This method is not as light as the previous work like Prompt-DT.\n- This paper shows the sub-performance compared with prompt DT when the demonstration set is expert.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. Explain more about the training hours and parameters for this method.**\n\nA1. The following tables show the number of model parameters, the training time, and inference time for one episode. Compared to DT-based baselines, our method introduces a lightweight world model that consists of about 15\\% parameters of the DT backbone. Thanks to this lightweight design, Meta-DT incurrs about 10\\% runtime increase during training and inference. **Considering the significant performance gain of our method, this lightweight cost is likely acceptable.** Moreover, Meta-DT's computation cost is even lower than other baselines.\n\n|`Model Size`|CMT|Prompt-DT|Generalized DT|CORRO|CSRO|Meta-DT|\n|---|---|---|---|---|---|---|\n|Point-Robot|675,989|603,259|600,754|705,961|858,480|702,758|\n|Cheetah-Vel|730,891|657,819|629,718|740,273|895,448|755,996|\n|Cheetah-Dir|730,891|657,819|629,718|740,273|895,448|755,996|\n|Ant-Dir|734,500|661,284|631,240|754,357|910,527|758,965|\n|Hopper-Param|726,079|653,199|627,651|756,827|885,014|751,274|\n|Walker-Param|729,720|658,872|629,286|735,665|893,808|755,129|\n\n|`Training time (s)`|CMT|Prompt-DT|Generalized DT|CORRO|CSRO|Meta-DT|\n|---|---|---|---|---|---|---|\n|Point-Robot|874|762|594|1,315|980|824|\n|Cheetah-Vel|1,068|758|600|1,609|1,076|1,050|\n|Cheetah-Dir|912|790|596|1,599|1,052|877|\n|Ant-Dir|5,076|4,605|3,021|9,880|6,179|4,930|\n|Hopper-Param|1,305|770|596|2,113|1,073|1,018|\n|Walker-Param|1,304|760|595|2,211|1,067|1,027|\n\n|`Inference time (ms)`|CMT|Prompt-DT|Generalized DT|CORRO|CSRO|Meta-DT|\n|---|---|---|---|---|---|---|\n|Point-Robot|9.52|8.68|46|10|15|9.16|\n|Cheetah-Vel|50.44|45.84|468|125|168|48.92|\n|Cheetah-Dir|50.40|48.2|467|123|167|48.90|\n|Ant-Dir|55.72|50.11|518|162|244|53.24|\n|Hopper-Param|51.77|48.64|488|146|210|51.00|\n|Walker-Param|54.13|49.60|516|147|231|53.01|\n\n**Q2. Explain the experiments why the method does not outperform the baselines in the expert demonstration.**\n\nA2. In experiments with Expert datasets as shown in Table 4 and Figure 10, Meta-DT outperforms three baselines (Generalized DT, CORRO, and CSRO) to a large extent in all environments. \n\nCompared to Prompt-DT, Meta-DT obtains **significantly better performance** in Point-Robot (-6.90 vs. -7.99), Cheetah-Vel (-52.42 vs. -133.78), and Ant-Dir (961.27 vs. 678.07), and obtains **extremely close performance** in Cheetah-Dir (874.91 vs. 960.32), Hopper-Param (383.51 vs. 393.79), and Walker-Param (437.79 vs. 449.15). **Taken together, it can be empirically verified that Meta-DT still outperforms Prompt-DT when evaluated on Expert datasets.**\n\nWith Expert datasets, Prompt-DT can access high-quality expert demonstrations as the task prompt to guide policy generation. Within the expert data, the behavior policy exactly equals the optimal policy and totally aligns with task characteristics, thus it can capture task-relevant information very well. Hence, Prompt-DT can achieve the same level of performance as Meta-DT in several environments. When faced with inferior datasets, the task information could be entangled with the features of behavior policies, producing biased task inference at test time. Hence, the performance of Prompt-DT might drop a lot, while Meta-DT can still obtain satisfactory performance due to accurately capturing task information via world model disentanglement.\n\n\n**Q3. This method is not as light as the previous work like Prompt-DT.**\n\nA3. We agree with you on this point, as our method introduces a pretrained context-aware world model. In contrast, Prompt-DT uses the same network architecture as DT, with some expert demonstrations prepended to the DT's input.\n\nOur motivation of introducing the additional world model is to disentangle task-relevant information from behavior policies, thus more accurately inferring task beliefs. With this effective disentanglement, our method is verified to be robust to the dataset quality and is more practical with fewer prerequisites in real-world scenarios. This is what Prompt-DT lacks, since it is sensitive to the quality of prompt demonstrations and can achieve satisfactory performance only when expert demonstrations are available at test time. \n\n\n### `Summary Response`\n\nThank you for your valuable review comments, which help us gain more insights on comparison to existing works like Prompt-DT, and further enhance our experimental illustration. We are honored to have **your recognition on our method, especially on its novelty** (\"a novel Meta-DT method\") **and superior performance** (\"shows the good performance... outperform the baselines\"). \n\nAlso, we are grateful that **most reviewers** (we refer to LRM6 as R2, EjPB as R3, sbh9 as R4, and BWEr as R5) **are quite affirmative to our overall contributions, including the novelty and motivation** (\"a clear motivation... quite novel and inspiring\" by R2, \"a new architecture... no major flaws with the paper\" by R4, and \"a new framework\" by R5), **the extensiveness of empirical evaluation** (\"comprehensive experiments on multiple benchmarks\" by R2, \"strong experimental evaluation over many different tasks\" by R4, and \"comprehensive experimental results\" by R5), and **the superior performance** (\"showing notable performance gains\" by R2, \"show the performance gain\" by R3, \"strong results in this evaluation\" by R4, and \"achieves superior generalization... demonstrating its practical effectiveness\" by R5). \n\nPlease let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HOY08fbw47", "reviewer": "Reviewer_LRM6", "summary": "This paper proposes a novel meta-RL framework called Meta Decision Transformer (Meta-DT). It leverages robust task representation learning via world model disentanglement to achieve task generalization in offline meta-RL. Firstly, it pretrains a context-ware world model to capture the task-relevant information from the offline dataset. Then,  it guides the sequence generation of decision transformer using the task representation and the self-guided prompt from past trajectories. The authors conduct extensive experiments on multiple benchmarks. Meta-DT demonstrates better few- and zero-shot generalization ability than other baselines.", "questions": "Please also consider responding to the Weaknesses.", "limitations": "The authors have discussed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper has a clear motivation in disentangling the task specific information. It is quite novel and inspiring to make use of the prompt to provide task-specific context and exploit architecture inductive bias.\n\n2. The authors conduct comprehensive experiments on multiple benchmarks, showing notable performance gains in many tasks of both few-shot and zero-shot scenarios. The ablation studies are also designed properly to justify the algorithm design.\n\n3. The paper is well-organized and easy to understand. The authors also provide detailed pseudo-codes and hyper-parameters in the appendix for reproduction.", "weaknesses": "1. From my perspective, the generalization ability of meta-DT depends heavily on the extrapolation ability of the context encoder. Although it is intuitive for some continuous tasks like Cheetah-Vel, this extrapolation tends to be quite hard for some other discrete tasks like Cheetah-Dir. Despite the good experiment performance, some intuitive explanation and theoretical justification are also desirable. \n\n2.  To evaluate the generalization ability more comprehensively, it may be helpful to provide the experiments with different number of training tasks. It is also interesting to see the extrapolation ability limit of meta-DT. For example, what if the training tasks of Ant-Dir are sampled from U[0,1.5$\\pi$] with hold-out tasks are from U[1.5$\\pi$, 2$\\pi$].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. Some intuitive explanation on the extrapolation ability are also desirable.**\n\nA1. Thank you for your insightful comments. Meta-DT's generalization comes from the extrapolation ability of the context-aware world model $W(r,s'|s,a; z^i)$. The world model completely describes the characteristics of the MDP/task. In the meta-learning setting, we assume that the world model shares some common structure across the task distribution, with a latent representation $z^i$ to approximate the unknown context of task $i$. **The principle behind is that tasks with similar contexts will behave similarly in the world model.** Intuitively, we map tasks into a latent space projected by the world model, and infer task similarity via the world model dynamics. That is, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. For discrete tasks like Cheetah-Dir, the world model also shares some common structure across the task distribution, e.g., the kinematics principle or locomotion skills. Hence, **the extrapolation of the world model works for both continuous and discrete tasks**. We will enhance the intuitive explanation of our method in the revised paper.\n\n\n**Q2. It may be helpful to provide the experiments with different numbers of training tasks.**\n\nA2. Following your advice, we have conducted new experiments to evaluate Meta-DT with a different number of training tasks. The following table shows the few-shot test performance of Meta-DT using Medium datasets (Figure S4 in Global Response). **Increasing the number of training tasks can usually improve the generalization ability to test tasks, especially in harder tasks like Ant-Dir**. \n\nThe result also matches the intuition of function approximation in machine learning. The \"generalization\" refers to the question: How can experience with a limited subset of the data space be usefully generalized to produce a good approximation over a much larger subset? In the single-task setting, increasing the valid sample points (before being overfitted) can generally improve the function approximation of the sample space and the generalization to unseen samples at testing. In the meta-learning setting, increasing the valid task points (before being overfitted) can usually boost the function approximation of the task space and the generalization to unseen tasks at testing. We will include this new experiment and the interesting findings in our revised paper.\n\n| Number of training tasks | 15 | 30 | 45 | \n| --- | --- | --- | --- |\n| Point-Robot | -10.20 $\\pm$ 0.28 | -10.45 $\\pm$ 0.22 | **-10.18** $\\pm$ 0.18 |\n| Cheetah-Vel | **-85.42** $\\pm$ 1.59 | -95.14 $\\pm$ 0.68 | -99.28 $\\pm$ 3.96 |\n|   Ant-Dir   | 388.50 $\\pm$ 9.02 | 402.71 $\\pm$ 5.00 | **412.00** $\\pm$ 11.53 |\n\n**Q3. It is also interesting to see the extrapolation ability limit of meta-DT.**\n\nA3. Following your suggestion, we have conducted new experiments on Ant-Dir to evaluate Meta-DT's generalization ability to out-of-distribution (OOD) tasks. The following table shows the few-shot performance of Meta-DT with OOD test tasks (Figure S5 in Global Response). In this case, we sample training tasks with a goal direction of $U[0, 1.5\\pi]$, and then test on tasks of $U[1.5\\pi, 2\\pi]$.\n\nObviously, **Meta-DT can still obtain better performance than baselines on OOD test tasks**, which again verifies our superiority. As stated in our A1 to Q1, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. For tasks like Ant-Dir, the world model shares some common structure across the task distribution (even for OOD tasks), e.g., the kinematics principle or locomotion skills. Hence, the extrapolation of the world model also works for OOD test tasks in this case. We will include this interesting experiment in the revised paper.\n\n| OOD Testing  | CMT | Prompt-DT | Generalized DT | CORRO | CSRO | Meta-DT |\n| --- | --- | --- | --- | --- | --- | --- |\n| Ant-Dir | 271.74 $\\pm$ 102.81 | 281.73 $\\pm$ 12.29 | 148.83 $\\pm$ 6.06 | 141.79 $\\pm$ 94.20 | 208.92 $\\pm$ 15.47 | **412.67** $\\pm$ 5.15 |\n\n### `Summary Response`\n\nThank you for your valuable suggestions, which help us gain more critical insights into the extrapolation ability of Meta-DT, and further enhance our experimental evaluation. The corresponding figures of extended experiments can be found in the PDF of Global Response. We are honored to have **your recognition of our method**, and are also grateful that **most reviewers are quite affirmative to our overall contributions**, including the novelty and motivation, the extensiveness of empirical evaluation, and the superior performance.\n\nPlease let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "owFCdtII29", "reviewer": "Reviewer_EjPB", "summary": "This work proposes a transformer-based framework for offline meta-reinforcement learning problems. The proposed algorithm utilizes a context-aware world model for task encoding and self-guided prompting. It outperforms existing offline meta-RL baselines.", "questions": "Q1. How do the world modeling and prompt tuning in the proposed method differ from CMT?\n\nQ2. How does the method perform in comparison to CMT?\n\nQ3. How does the method perform in $\\texttt{Humanoid-Dir}$ and more meta-world benchmarks such as  $\\texttt{Hammer}$, $\\texttt{Sweep}$, and $\\texttt{Door-Lock}$?\n\nQ4. What is the sensitivity with the prompt length $k$?", "limitations": "The limitations section is inadequate. \n\nL1. There seems to be high computational overhead during training and inference.\n\nL2. Meta-world benchmarks concluded only from one task which may question the generalization of the proposed method.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "S1. The results in the paper show the performance gain of the proposed method over baselines.", "weaknesses": "W1. The work misses the CMT baseline which is cited but not shown as a baseline (https://arxiv.org/abs/2211.08016).\n\nW2. The work lacks novelty as the idea of a world model and Prompt tuning already exists in CMT. Moreover, the transformer architecture is based on a decision transformer.\n\nW3. The run-time is unclear for fare comparison with baselines.\n\nW4. The results for Mujoco benchmark $\\texttt{Humanoid-Dir}$ in CSRO are missing. Meta-world results are limited (shown results only for $\\texttt{Reach}$), thus hard for a general conclusion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. The work misses the CMT baseline which is cited but not shown as a baseline.**\n\nA1. The reason why we did not include CMT as a baseline is that **its code has not been open-sourced**. We emailed CMT's authors for requesting the source code, and received the response that the code would be made public at an uncertain time. Following your suggestion, **we have included CMT as a new baseline** using our own implementation (Figure S1 of Global Response). **The results again demonstrate the consistent superiority of Meta-DT over CMT**.\n\n|Method|Point-Robot|Cheetah-Vel|Cheetah-Dir|Ant-Dir|Hopper-Param|Walker-Param|Reach| \n|---|---|---|---|---|---|---|---|\n|CMT|-12.21$\\pm$ 0.51|-131.71$\\pm$ 10.51|545.55$\\pm$ 7.07|275.67$\\pm$ 19.12|325.20$\\pm$ 3.62|356.45$\\pm$ 5.58|2295.34$\\pm$ 121.56|\n|Meta-DT|**-10.18**$\\pm$ 0.18|**-99.28**$\\pm$ 3.96|**608.18**$\\pm$ 4.18|**412.00**$\\pm$ 11.53|**348.20**$\\pm$ 3.21|**405.12**$\\pm$ 11.11|**2458.83**$\\pm$ 162.10|\n\n**Q2. How do the world modeling and prompt tuning in the proposed method differ from CMT?**\n\nA2. **We do not tune the prompt**. Meta-DT selects the trajectory segment using the world model to construct the prompt, without tuning any model parameters associated with the prompt. While both involve the concepts of world model and prompt, **Meta-DT differs significantly from CMT in many crucial aspects** including:\n\n- `The way of using the world model for task inference is different.` CMT infers task beliefs at the trajectory level, abstracting a task prompt $z_{\\tau}$ from a context trajectory $\\tau$ and using that prompt to guide policy generation for the entire episode. In contrast, Meta-DT performs task inference at the transition level, abstracting task representation $z_t$ at each timestep $t$. Hence, Meta-DT can provide more fine-grained guidance to realize high-capacity generalization, and can achieve zero-shot generalization via deriving real-time task inference.\n\n- `The prompt design is quite different.` CMT abstracts task presentation $z_{\\tau}$ from a context trajectory and uses $z_{\\tau}$ as the task prompt. In contrast, we directly use a trajectory segment $\\tau^*$ as the task prompt, enjoying the power of architecture inductive bias like in Prompt-DT. Moreover, CMT tunes the prompt adaptor layer based on relabeling the offline dataset, while Meta-DT does not involve prompt tuning. Our complimentary prompt design is more lightweight and easy-to-implement. By using the world model to help construct the prompt, the world model (algorithmic perspective) and the complementary prompt (architecture perspective) work as an efficient whole to perform accurate task inference for Meta-DT.\n\n- `The pipeline of training the world model and meta-policy is different.` CMT simultaneously learns the world model and meta-policy. In contrast, Meta-DT decouples the learning process, which pretrains the world model and then trains the meta-policy with a fixed world model. This kind of decoupling assures a more stable process for learning task representations, since the algorithm can solely focus on capturing task similarity via the world model dynamics. \n\n- `The algorithm paradigm is different.` Our method is based on decision transformer, a return-conditioned supervised learning paradigm that has achieved promising results on offline RL. It feeds the (return-to-go, state, action) tuple sequence to the causal transformer to predict the action element only. In contrast, CMT uses another paradigm that feeds the (state, action, reward) tuple sequence to the causal transformer to perform autoregressive prediction on all elements.\n\n**Q3. Evaluation on Humanoid-Dir and more Meta-World environments.**\n\nA3. Following your suggestion, **we have conducted new evaluation on Humanoid-Dir and two Meta-World environments of Sweep and Door-Lock**. The following table shows the few-shot testing performance of Meta-DT against extended baselines using Medium datasets (Figure S2 in Global Response). The results again demonstrate the **consistent superiority of Meta-DT** over all baselines in these domains.\n\n|Environment|CMT|Prompt-DT|Generalized DT|CORRO|CSRO|Meta-DT|\n|---|---|---|---|---|---|---|\n|Humanoid-Dir|548.42$\\pm$ 40.98|556.43$\\pm$ 15.84|538.95$\\pm$ 7.04 |429.66$\\pm$ 150.58|465.58$\\pm$ 21.78|**661.08**$\\pm$ 57.04|\n|Sweep|493.06$\\pm$ 72.54|593.11$\\pm$ 92.10|549.17$\\pm$ 79.35|272.67$\\pm$ 77.41|446.50$\\pm$ 191.46|**760.18**$\\pm$ 4.87|\n|Door-Lock|2291.82$\\pm$ 76.88|2284.95$\\pm$ 341.20|2601.24$\\pm$ 44.42 |1615.66$\\pm$ 248.00|2216.67$\\pm$ 116.84|**3025.92**$\\pm$ 23.41|\n\n**Q4. The run-time during training and inference is unclear for fare comparison with baselines.**\n\nA4. Please refer to our `A1 to Reviewer iZYy` for the tables of model parameters, training time, and inference time for one episode. Compared to DT-based baselines, our methods incurrs about 15\\% increase in model parameters, and about 10\\% increase in training and inference time. **Considering the significant performance gain of our method, this lightweight cost is likely acceptable.** Moreover, Meta-DT's computation cost is even lower than other baselines.\n\n**Q5. What is the sensitivity with the prompt length $k$?**\n\nA5. We have conducted new experiments to investigate the influence of prompt length $k$ on Meta-DT’s performance (Figure S3 in Global Response). Generally, **the performance of Meta-DT is not sensitive to the pre-defined value of prompt length**.\n\n|Prompt length $k$|3|5|7| \n|---|---|---|---|\n|Point-Robot|**-10.18**$\\pm$ 0.18|-10.19$\\pm$ 0.15|-10.22$\\pm$ 0.04|\n|Cheetah-Dir|606.69$\\pm$ 2.92|**608.18**$\\pm$ 4.18|605.09$\\pm$ 6.31|\n|Ant-Dir|408.97$\\pm$ 18.01|412.00$\\pm$ 11.53|**413.09**$\\pm$ 8.82|\n\n**Q6. The limitations section is inadequate.**\n\nA6. We have discussed three threads of our limitations. During rebuttal, the valuable comments from all reviewers also give us new insights into our method, and we will have a more thorough discussion on limitations based on these comments.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nYzf3BuFsG", "reviewer": "Reviewer_sbh9", "summary": "The paper introduces a new architecture for Meta-RL based on Decision Transformers. The new architecture uses a world model responsible for efficiently encoding task information from demonstrations. The model also introduces a prompt encoder which acts as a boosting mechanism to the context encoder to enhance the models task adaptation abilities.", "questions": "Why do you thin the disentanglement works well as an idea? I understand the different perspectives of  adaptation the prompt vs context encoder aim to target, but could you not form the context in manner to do this, for example running it twice with different subsets of experience? Is it the way things are fed in to the prompt vs context encoder (segmenting) and the lack of constraints on the prompt encoders representation?", "limitations": "The limitations and future work are well addressed", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Strong experimental evaluation including over many different tasks, settings,f ablations, and comparison against several other techniques\n- Strong results in this evaluation, with good demonstration that the introduced components make a difference.\n- Simple design with good results\n- Good figures and explanations, with std error ranges -- which can be critical for RL", "weaknesses": "No major flaws with the paper.\nThe evaluation covers environments which are relatively homogenous in nature. Performance in MetaWorld is much less differentiating compared to MuJoCo. It would be nice to see more detailed evaluation in non-mujoco tasks. \n\nIt would be interesting to see how the technique generalizes across worlds and tasks with more diversity in them, mujoco worlds are extremely minimal with little variation. Something like training on K-levels of an Atari game and generalizing to the remaining N-K levels of well suited Atari games. It is entirely possible the research in this simply isn't ready to tackle hard cases such as this though. Or something like the worlds seen in the Muesli paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. It would be nice to see more detailed evaluation in non-mujoco tasks.**\n\nA1. Thank you for your advice. We have conducted new evaluations on Humanoid-Dir and two more Meta-World environments of Sweep and Door-Lock. The following table shows the few-shot testing performance of Meta-DT against baselines using Medium datasets (Figure S2 in Global Response). **The results again demonstrate the consistent superiority of Meta-DT over all baselines in these newly evaluated domains**.\n\n| Environment  | CMT | Prompt-DT | Generalized DT | CORRO | CSRO | Meta-DT |\n| --- | --- | --- | --- | --- | --- | --- |\n| Humanoid-Dir | 548.42 $\\pm$ 40.98 | 556.43 $\\pm$ 15.84 | 538.95 $\\pm$ 7.04 | 429.66 $\\pm$ 150.58 | 465.58 $\\pm$ 21.78 | **661.08** $\\pm$ 57.04 |\n|     Sweep    | 493.06 $\\pm$ 72.54 | 593.11 $\\pm$ 92.10 | 549.17 $\\pm$ 79.35 | 272.67 $\\pm$ 77.41 | 446.50 $\\pm$ 191.46 | **760.18** $\\pm$ 4.87 |\n|  Door-Lock   | 2291.82 $\\pm$ 76.88 | 2284.95 $\\pm$ 341.20 | 2601.24 $\\pm$ 44.42 | 1615.66 $\\pm$ 248.00 | 2216.67 $\\pm$ 116.84 | **3025.92** $\\pm$ 23.41 |\n\n\n**Q2. It would be interesting to see how the technique generalizes across worlds and tasks with more diversity in them, like Atari games.**\n\nA2. Thank you for pointing out this insightful direction. As stated in the limitations section, our generalist model is trained on relatively lightweight datasets compared to popular large models. The urgent trend that RL practitioners are striving to break through is to deploy on significantly large datasets with more diversity, unlocking the scaling law with the transformer architecture. Due to limited time during rebuttal, it is hard for us to conduct systematic empirical evaluation on Atari domains. We will include the above insights into our limitations section, and are looking forward to investigating this promising future direction.\n\n\n**Q3. Why do you think the disentanglement works well as an idea?**\n\nA3. Thank you for your insightful question. Existing context-based or prompt-based methods usually infer task representations by feeding subsets of experience into a trajectory encoder. The offline dataset depends on both the task and the behavior policy. **The task information could be entangled with the features of behavior policies**, thus producing biased task inference at test time due to the shift in behavior policies.\n\nHence, we attempt to find a stable way to accurately disentangle task-relevant information from behavior policies. In the RL regime, **the world model completely describes the characteristics of the MDP/task, and keeps invariant to behavior policies or collected datasets**. Naturally, it could be a promising alternative for accurately disentangling task beliefs. That is the beginning of motivating the world model disentanglement as our idea. \n\nIn the meta-learning setting, we assume that the world model shares some common structure across the task distribution. **The principle behind is that tasks with similar contexts will behave similarly in the world model.** Intuitively, we map tasks into a latent space projected by the world model, and infer task similarity via the world model dynamics. That is, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. Subsequently, we pretrain the world model and then train the meta-policy with a fixed world model. This kind of decoupling assures a more stable process for learning task representations, since the algorithm can solely focus on capturing task similarity via the world model dynamics. \n\n### `Summary Response`\n\nThank you for your valuable comments, which help us enhance our experimental evaluation and gain more insightful directions for future work. We are honored to have **your recognition of our work**, and are also grateful that **most reviewers are quite affirmative to our overall contributions**, including the novelty and motivation, the extensiveness of empirical evaluation, and the superior performance.\n\nThe corresponding figures of extended experiments can be found in the PDF of Global Response. Please let us know if we have addressed your concerns, and we are more than delighted to have further discussions and improve our manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZQCpGPIGm0", "reviewer": "Reviewer_BWEr", "summary": "A new framework is proposed for offline meta-reinforcement learning (meta-RL) that leverages transformers and world model disentanglement to enhance task generalization without the need for expert demonstrations or domain knowledge. The approach, called Meta Decision Transformer (Meta-DT), utilizes a context-aware world model to encode task-relevant information which guides the transformer in generating task-oriented sequences. Experimental results show that Meta-DT achieves superior few-shot and zero-shot generalization across different benchmarks, indicating its practicality and effectiveness in RL applications.", "questions": "Mentioned in the weakness section", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Generalization without Experts: Meta-DT effectively generalizes across unseen tasks without requiring expert demonstrations or domain knowledge, which is a significant advancement in the field of meta-reinforcement learning.\n\nWorld Model Disentanglement: The introduction of a context-aware world model that disentangles task-relevant information from behavior policies enhances the robustness and accuracy of task representation.\n\nEfficient Use of Transformers: The paper creatively applies the transformer architecture to offline meta-RL, leveraging its strong sequential modeling capabilities to improve policy generation and adaptation.\n\nEmpirical Validation: The paper includes comprehensive experimental results showing that Meta-DT outperforms existing baselines in few-shot and zero-shot generalization tasks across multiple benchmarks, demonstrating its practical effectiveness.", "weaknesses": "1. Are there any environments or types of tasks where Meta-DT's performance is notably limited? What are the challenges in extending the model to such environments, and how might these be addressed in future work?\n\n2. Can you elaborate on the process and benefits of world model disentanglement in Meta-DT? How does this process affect the model's ability to generalize across different tasks, especially in dynamically changing environments?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q1. Are there any environments or types of tasks where Meta-DT's performance is notably limited? What are the challenges in extending the model to such environments, and how might these be addressed in future work?**\n\nA1. Thank you for your insightful questions. We evaluated Meta-DT on seven environments adopted from three classical benchmarks in meta-RL. During rebuttal, we also conducted experiments on new evaluation domains, including Humanoid-Dir and two more Meta-World environments (Figure S2 in Global Response). **Results showed the consistent superiority of Meta-DT in these ten environments.** \n\nOn the other hand, the above environments are relatively homogenous in nature. As stated in the limitations section, our generalist model is trained on relatively lightweight datasets compared to popular large models. The urgent trend that RL practitioners are striving to break through is to deploy on significantly large datasets with more diversity. As mentioned by Reviewer sbh9, **it would be interesting to see how Meta-DT generalizes across worlds and tasks with more diversity**, like training on $K$-levels of an Atari game and generalizing to the remaining $N-K$ levels of well-suited Atari games. \n\nThe **challenges** in extending to such paradigms might include: i) how to collect large-scale datasets with sufficient diversity and high quality, ii) how to tackle the heterogenous task diversity in hard cases, and iii) how to scale RL models to very large network architectures like in large language or visual models. To tackle these potential challenges, we conjecture several **promising solutions** including: i) leveraging self-supervised learning to facilitate task representation learning at scale, ii) enabling efficient prompt design or prompt tuning to facilitate in-context RL, and iii) incorporating effective network architectures like mixture-of-experts to handle heterogenous domain diversity. We will include the above insights into our limitations section, and are looking forward to investigating these promising future directions.\n\n\n**Q2. Elaborate on the process and benefits of world model disentanglement in Meta-DT, and how it affects the model's generalization ability.**\n\nA2. Existing context-based or prompt-based methods usually infer task representations by feeding subsets of experience into a trajectory encoder. However, in RL regimes, since the offline dataset depends on both the task and the behavior policy, **the task information could be entangled with the features of behavior policies**, thus producing biased task inference at test time due to the shift in behavior policies. Taking 2D navigation as an example, tasks differ in goals and the behavior policy is going towards the goal for each task. The algorithm can easily distinguish tasks based on state-action distributions rather than reward functions, leading to extrapolating errors when the behavior policy shifts to random exploration during testing.\n\nHence, we attempt to find a stable way to accurately disentangle task-relevant information from behavior policies. For RL paradigms, **the world model completely describes the characteristics of the MDP/task, and keeps invariant to behavior policies or collected datasets**. Naturally, it could be a promising alternative for accurately disentangling task beliefs. That is exactly the motivation of our world model disentanglement. \n\nIn the meta-learning setting, we assume that the world model shares some common structure across the task distribution. **The principle behind is that tasks with similar contexts will behave similarly in the world model.** Intuitively, we map tasks into a latent space projected by the world model, and infer task similarity via the world model dynamics. That is, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. Subsequently, we pretrain the world model and then train the meta-policy with a fixed world model. This kind of decoupling assures a more stable process for learning task representations, since the algorithm can solely focus on capturing task similarity via the world model dynamics. \n\n\n### `Summary Response`\n\nThank you for your valuable comments, which help us gain more insights into our methodology and motivate promising directions for future work. We are honored to have **your recognition of our work**, and are also grateful that **most reviewers are quite affirmative to our overall contributions**, including the novelty and motivation, the extensiveness of empirical evaluation, and the superior performance.\n\nThe corresponding figures of extended experiments can be found in the PDF of Global Response. Please let us know if we have addressed your concerns, and we are more than delighted to have further discussions and improve our manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Revision Summary

We thank the reviewers for their valuable feedback (we refer to iZYy as R1, LRM6 as R2, EjPB as R3, sbh9 as R4, and BWEr as R5). We are grateful that **most reviewers are quite affirmative to our overall contributions**, including **the novelty and motivation** ("a novel Meta-DT" by R1, "a clear motivation... quite novel and inspiring” by R2, "a new architecture... no major flaws with the paper" by R4, and "a new framework" by R5), **the extensiveness of empirical evaluation** ("comprehensive experiments on multiple benchmarks" by R2, "strong experimental evaluation over many different tasks" by R4, and "comprehensive experimental results" by R5), and **the superior performance** ("shows the good performance" by R1, "showing notable performance gains" by R2, "show the performance gain" by R3, "strong results in this evaluation" by R4, and "achieves superior generalization... demonstrating its practical effectiveness" by R5).

We have made **a number of changes** to address all reviewers' suggestions and concerns. A short summary of the modifications is made as

1. We include CMT as a new baseline for experimental evaluation.
2. We conduct new empirical evaluation on Humanoid-Dir and two Meta-World environments of Sweep and Door-Lock.
3. We conduct new hyperparameter analysis experiments to investigate the influence of prompt length $k$ on Meta-DT’s performance.
4. We conduct new experiments to demonstrate Meta-DT's performance with a different number of training tasks. 
5. We include new experiments to evaluate Meta-DT's generalization ability to out-of-distribution tasks.
6. We show and analyze the number of model parameters, the training time, and the inference time for Meta-DT and baselines.
7. We further justify our methodology including: i) the significant differences from CMT and corresponding superiorities over CMT, ii) the motivation and benefits of world model disentanglement, and iii) limitations of our method, challenges in extending to harder cases, and promising future directions.

In summary, we have **significantly extended our empirical evaluation** based on the comprehensive experimental results in our original manuscript. Impressively, **the results of massively extended experiments are generally consistent with observations and conclusions from our original manuscript**. The tabular results are given in responses to each reviewer, and corresponding figures can be found in the PDF of Global Response.

Please let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.

Meta-review: All reviewers were positive about the paper after the rebuttal. Some reviewers especially praised the paper's simplicity and strong generalization results. AC liked the experiments with varying number of tasks which helped reveal the method's potential and limitation in generalization.
Decision: Accept (poster)