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_cc_2024_Conference
Published: 2024-04-23 09:38:20

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.

Main paper: 
 introduction: 
 Building generalist models that solve various tasks by training on vast task-agnostic datasets has emerged as a dominant paradigm in the language [1] and vision [2] communities. Offline reinforcement learning (RL) [3] 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 [4, 5, 6] , akin to the rise of large language models (LLMs) like GPT [7] . A performant example is the decision transformer (DT) [8] that leverages the transformer's strong sequence modeling capability for trajectory data and achieves promising results on offline RL [9, 10, 11, 12] .Transformer-based large models have shown remarkable scaling properties and high transferability across various domains, including language modeling [13, 7] , computer vision [14, 15] , and image generation [16, 17] . However, their counterparts in RL are usually specialized to a narrowly defined task and struggle with poor generalization to unseen tasks, due to distribution shift and the lack of Inspired by this, we propose a novel framework for offline meta-RL, named Meta-DT that leverages robust task representation learning via world model disentanglement to conduct task-oriented sequence 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 encoder 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 a 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-oriented context. Finally, we leverage the past trajectories generated by the meta-policy as a self-guided prompt to exploit the architecture inductive bias, akin to Prompt-DT [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 to encode 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. 
related work: 
 Offline Meta-RL. Offline RL [33] 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-based meta-learning methods [37, 38] seek policy parameter initialization that requires only a few adaptation steps to new tasks. MACAW [39] introduces this architecture into value-based RL subroutines, and uses simple supervised regression objectives for sample-efficient offline meta-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 [42] , CORRO [24] , CSRO [25] , and UNICORN [43] . These works extended from the online meta-RL setting are mainly trained by temporal difference learning, the dominant paradigm in RL [9, 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-RL 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 [7] , ViT [14] ,and DiT [46] . Decision transformer [8] first introduces the transformer's sequence modeling capacity to solving offline RL without temporal difference learning, which autoregressively outputs optimal actions conditioned on desired returns. As a concurrent study, trajectory transformer [47] directly models distributions over sequences of states, actions, and rewards, followed by beam search as a planning algorithm in a model-based manner. Extending the LLM-like structure to RL, this paradigm activates a new pathway toward scaling powerful RL engines with large-scale compute and data [48] .For multi-task learning, Gato [49] trains a multi-modal generalist policy in the GPT style with a decoder-only transformer. Multi-game DT [4] trains a suite of up to 46 Atari games simultaneously, and controls policy generation at inference time with a pretrained classifier that indicates the expert level of behaviors. For generalization to unseen tasks, Prompt-DT [21] exploits a prefix prompt architecture and Generalized DT [22] designs a 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 3 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. 
offline meta reinforcement learning: 
 RL is generally formulated as a Markov decision process (MDP) M=⟨S, A, T, R, γ⟩, where S/A is the state/action space, T /R is the state transition/reward function, and γ is the discount. The objective is to find a policy π(a|s) to maximize the expected return asmax π J M (π) = E π [ ∞ t=0 γ t R(s t , a t )].In offline meta-RL, we assume that the task follows a distribution M i = ⟨S, A, T i , R i , γ⟩ ∼ P (M ), where tasks share the same state-action spaces while vary in the reward and state transition functions, i.e., environment dynamics. For each task i from a total of N training tasks {M i } N i=1 , an offline dataset D i = {(s i,j , a i,j , r i,j , s ′ i,j )} J j=1 is collected by an arbitrary behavior policy π i β . The agent can only access the offline datasets {D i } N i=1 to train a meta-policy π meta . At test time with the few-shot setting, given an unseen task M ∼ P (M ), the agent can access a small dataset D = {(s j , a j , r j , s ′ j )} J ′ j=1 to construct the task prompt or infer the task belief before policy evaluation. For the zero-shot setting, the trained meta-policy is directly evaluated on the unseen task by interacting with the test environment to estimate the expected return. The objective is to learn a meta-policy to maximize the expected episodic return over test tasks asJ(π meta ) = E M ∼P (M ) [J M (π meta )].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 [8] introduces associated advances to solve RL problems and re-frames offline RL as a return-conditioned sequential modeling problem, inheriting the transformer's efficiency and scalability when modeling long trajectory data. DT models the probability of the next sequence token x t conditioned on all tokens prior to it: P θ (x t |x <t ), similar to contemporary decoder-only sequence models [51] . The sequence we consider has the form:x = (• • • , Rt , s t , a t , • • • ),where Rt is the agent'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, Rt = T i=t r i , and during testing, Rt = G * -t i=0 r i , where where G * is the target return for an entire episode. The same timestep embedding is concatenated to the embeddings of Rt , s t , and a t , and the head corresponding to the state token is trained to predict the action by minimizing its error from the ground truth. t , and the generalized decoders (R ϕ , T φ ) that predict the reward and next state conditioned on z i t . Then, the inferred task representation is injected as a contextual condition to the causal transformer to guide task-oriented sequence generation. Finally, we design a self-guided prompt from history trajectories generated by the meta-policy at test time. We select the trajectory segment that yields the largest prediction error on the pretrained world model, aiming to encode task-relevant information complementary to the world model maximally. 
method: 
 In this section, we present Meta Decision Transformer (Meta-DT), a novel offline meta-RL 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. 1 illustrates the method overview, and the following subsections present the detailed implementations.The key to efficient generalization is how to accurately encode task-relevant information for extrapolating knowledge across tasks, which proves more challenging in RL. A supervised or unsupervised 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 an 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.To address this intrinsic challenge, we ought to disentangle task-relevant information from behavior policies. The transition dynamics, i.e., the reward and state transition functions p(s ′ , r|s, a), completely describes the characteristics of the underlying environment. Naturally, this environment dynamics, also called world model, keeps invariant to behavior policies or collected datasets, and could be a promising alternative for accurately inferring task beliefs. Therefore, we 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 each 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, we use a latent representation z to approximate that variable. For a given task M i , the reward and state transition functions can be approximated by a generalized context-aware world model W that is shared across tasks asEQUATIONSince we do not have access to the true task description or identity, we need to infer task representation z i t given the agent's h-step experience up to timestep t collected in task M i asµ i t = (s t-h , a t-h , r t-h , ..., s t-1 , a t-1 , r t-1 , s t ).(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 dynamics-specific information into a latent task representation, and ii) decoding the environment dynamics conditioned on that representation. First, we use a simple yet effective context encoder E ψ to embed recent experiences into a task representation as z i t = E ψ (µ i t ). Second, the decoder contains a generalized reward model R ϕ and state transition model T φ . The task representation is augmented into the input of both models to predict the instant reward rt+1 = R ϕ (s t , a t ; z i t ) and next state ŝt+1 = T φ (s t , a t ; z i t ). 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 taskrelevant information from training environments. The context encoder is jointly trained by minimizing the reward and state transition prediction error conditioned on the task representation asL(ψ, ϕ, φ) = E µ i t ∼Di E z i t ∼E ψ (µ i t ) r t+1 -R ϕ (s t , a t ; z i t ) 2 + s t+1 -T φ (s t , a t ; z i t ) 2 . (3)Recent works suggest the prompt framework as an effective paradigm for pretraining transformerbased large models on massive datasets and adapting them to new scenarios with few or no labeled data [54] . Prompt-DT [21] adopts that paradigm to RL problems by prepending a task prompt to the DT's input. At test time, the agent is assumed to access a handful of expert demonstrations to construct the prompt and perform few-shot generalization to new tasks. However, in real-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 an initially unknown environment. Recent studies [36, 23] also suggest that the quality of demonstrations must be high enough to act as a 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 a 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-quality 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 perspective) 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 a sequence segment containing multiple tuples, ( R * , s * , a * ), which are sampled from the agent's history trajectories. Each element with superscript * is associated with the prompt. For a given task M i , we obtain a k-step prompt τ * i asτ * i = R * j , s * j , a * j , R * j+1 , s * j+1 , a * j+1 , ..., R * j+k , s * j+k , a * j+k ,where j = maxEQUATION)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 can store partial information to implicitly capture task dynamics.Our Meta-DT architecture is built on decision transformer [8] 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 it fixed to produce compact task representations for the downstream meta-training process. For each task M i , we use the context encoder from the pretrained world model to infer the contextual information for each timestep as z i t = E ψ (µ i t ). Then, the input of Meta-DT consists of two parts: i) the k-step complementary prompt τ * i derived from (4), and ii) the most recent K-step history τ + i augmented by the learned task representations asEQUATIONThe input sequence (τ * i , τ + i ) 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-step history. At test time with a few-shot setting, the meta-policy is allowed to interact with the environment for a 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-DT 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.1 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 Meta-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 [25] ; ii) the multi-task MuJoCo control [55, 36] , containing Cheetah-Vel, Cheetah-Dir, 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 a distribution of tasks and divide them into the training set T train and test set T test . On each training task, we use SAC [57] to train a single-task policy independently for collecting datasets. We consider three ways to construct offline datasets: Medium, Mixed, and Expert. More details about environments and datasets are given in Appendix B.Baselines. We compare Meta-DT to four competitive baselines that cover two distinctive paradigms in offline meta-RL: the DT-based 1) Prompt-DT [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 5 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.Note that these baselines work under the few-shot setting, since they require expert demonstrations as task prompts or some warm-start data to 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 1 present the testing curves and converged performance of Meta-DT and baselines on various domains using Medium datasets under an aligned few-shot setting. In these environments with varying reward functions or state transition dynamics, Meta-DT consistently obtains superior performance regarding data efficiency and final asymptotic 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.Since Meta-DT can derive real-time task representations z t from its h-step experience µ t via the pretrained context encoder, it 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 Table 2 present 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 a collapsed performance (about 20%-35% drop on average) compared to their fewshot counterparts. The main reason is that they usually require expert demonstrations as task prompts or some high-quality warmup data to accurately capture task information. In contrast, Meta-DT can abstract compact task presentations via the context-aware world model, and its performance in zero-shot scenarios drops merely a little (about 5% on average) compared to the 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-DT with fewer prerequisites in real-world applications. To test the respective contributions of each component, we compare Meta-DT to four ablations: i) w/o_context, it only removes the task representation z i t from the input of the casual transformer; ii) w/o_com, it randomly chooses a segment from a candidate trajectory to construct the prompt, i.e., removing the complementary way for constructing the prompt; iii) w/o_prompt, it directly removes the prompt component; iv) DT, it removes all components and degrades to the original DT. For ablations, all other structural modules are kept consistent strictly with the full Meta-DT.Fig. 4 and Table 3 show ablation results on representative domains using Medium datasets. First, Meta-DT obtains a decreased performance when any component is removed, which demonstrates that all components are 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 a more vital role in capturing task-relevant information. Third, the performance order of w/o_prompt < w/o_com < 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 is more significant in complex environments like Ant-Dir than in simpler ones like Point-Robot. It again verifies our superiority in tackling challenging tasks.A good offline algorithm ought to be robust to different types of datasets that involve a wide range of behavioral policies [58] . To test how Meta-DT performs with data of different qualities, we also conduct experiments on the Mixed and Expert datasets. Meta-DT and baselines with a few-shot setting, and Table 4 shows corresponding converged results.In Mixed datasets, Meta-DT still outperforms all baselines by a large margin, especially in 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 Point-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 a lot when provided with prompts from Medium or Mixed datasets, also mentioned in the original paper [21] 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 a significant performance gap compared to Meta-DT. In summary, the above results verify that our method is robust to the dataset quality and is more practical with fewer prerequisites in real-world scenarios. In the paper, we tackle the offline meta-RL challenge via leveraging advances in sequence modeling with scalable transformers, marking a 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 of pre-training the world model and training the causal transformer. A promising improvement is to develop a unified framework that simultaneously abstracts the task representation and learns the meta-policy, akin to in-context learning [59] . Also, our generalist model is trained on relatively lightweight datasets compared to popular large models. A 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 are striving to break through. Another interesting line is to leverage self-supervised learning [15, 60] to facilitate task representation learning at scale. We leave these as future work.Compute the predicted reward rt = R ϕ (s t , a t ; z i t ) and next state ŝt+1 = T φ (s t , a t ; z i t )9 end 10 Update E ψ , R ϕ , and T φ using the loss as Infer the context of each timestep t from its h-step history using the pretrained context encoder as z t = E ψ (s t-h , a t-h , r t-h , ..., s t-1 , a t-1 , r t-1 , s t )EQUATION9Augment the trajectory τ i with per-step context as Augment the trajectory τ i with context asEQUATIONa pred = F θ (τ input i ), ∀τ input i ∈ B 15 L(θ) = 1 B τ input i ∈B (a -a pred )EQUATIONStep env. and get feedback s, a, r, R ← R -r Infer the context from h-step history using E ψ as z t = E ψ (s t-H+1 , a t-H+1 , r t-H+1 , ..., s t-1 , a t-1 , r t-1 , s t ) 5Augment the trajectory τ i with context asEQUATIONGet action a = F θ (τ + i )• Mixed: a mixed dataset of the Medium and Expert types. We use a combination of the medium (70%) and expert (30%) datasets collected as above to generate the mixed dataset.Specifically, we generate 100 trajectories for the Point-Robot and MuJoCo environments, and 300 trajectories for the Meta-World environments, using the saved checkpoint policies to construct corresponding offline datasets.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 on 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, 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. • CORRO [24] , is a temporal difference-based 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 as 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 a context trajectory with an arbitrary policy to infer the task representation.• CSRO [25] , is a temporal difference-based method that proposes a 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-RL 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 to 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.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-DT, we ablate the prompt component and derive real-time task representations from the history trajectory via the pre-trained context encoder. For a fair comparison, we modify the baselines to 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). Context horizon h. As shown by the ablation studies in Sec. 5.3, task representation learning via the world model disentanglement plays a more vital role in capturing task-relevant information in Meta-DT. The task representation z i t is abstracted from the agent's h-step experience µ i t = (s t-h , a t-h , r t-h , ..., s t-1 , a t-1 , r t-1 , s t ). 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 h on Meta-DT's performance.  
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- This paper shows the sub-performance compared with prompt DT when the demonstration set is expert.
questions: questions: - Can you explain more about the training hours and parameters for this method?
- Can you explain the experiments why the method does not outperform the baselines in the expert demonstration?
limitations: limitations: See weaknesses and questions.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

2.  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$].
questions: questions: Please also consider responding to the Weaknesses.
limitations: limitations: The authors have discussed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: W1. The work misses the CMT baseline which is cited but not shown as a baseline (https://arxiv.org/abs/2211.08016).

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

W3. The run-time is unclear for fare comparison with baselines.

W4. 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.
questions: questions: Q1. How do the world modeling and prompt tuning in the proposed method differ from CMT?

Q2. How does the method perform in comparison to CMT?

Q3. How does the method perform in $\texttt{Humanoid-Dir}$ and more meta-world benchmarks such as  $\texttt{Hammer}$, $\texttt{Sweep}$, and $\texttt{Door-Lock}$?

Q4. What is the sensitivity with the prompt length $k$?
limitations: limitations: The limitations section is inadequate. 

L1. There seems to be high computational overhead during training and inference.

L2. Meta-world benchmarks concluded only from one task which may question the generalization of the proposed method.
paper_summary: paper_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.
strength_weakness: strength_weakness: No major flaws with the paper.
The 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. 

It 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.
questions: 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: limitations: The limitations and future work are well addressed
paper_summary: paper_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.
strength_weakness: strength_weakness: 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?

2. 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?
questions: questions: Mentioned in the weakness section
limitations: limitations: Yes

Comments: title: Thank you!
Looking forward to further discussions!
Thank you!
Thank you!
Response
Thank you!
Response Summary by Authors
comment: Thank you very much for your constructive feedback and your time in helping us improve our work. We sincerely appreciate you raising the score on our work!
Thanks for the detailed explanation and additional experiments. I will raise my score.
Dear Reviewer,

Thank you for your insightful comments. We were wondering if our response and revision have resolved your concerns. We have addressed your initial questions through our rebuttal and are eager to clarify any further points you might raise. Please feel free to provide additional feedback. We greatly appreciate your continued engagement.

Best regards,

Authors
Thank you for your time in helping us improve our work. We are more than delighted to see that we could address your all concerns. We sincerely appreciate you raising the score on our work!
Thank you for your time in helping us improve our work. We are happy to see that we could address your main concerns. Thank you sincerely for raising the score on our work. We truly appreciate it!
Thank you for the detailed response. As the majority of my concerns have been addressed, I will increase my rating.
All of my concerns are well addressed. Thank you. I will increase my rating to accept
Thank you for your time in helping us improve our work. We are happy to see that we could address your concerns. We sincerely appreciate your recognition of our contribution and your vote to accept our work!
Thank you for the detailed response. This has helped to clarify my questions. I recognize the contribution of this work, and I would like to keep my original rating to vote for acceptance.
Thank you for your valuable review comments, which help us gain more critical insights on the difference from existing works, and further enhance our empirical evaluation. We summarize the three main concerns as 

- `The difference from CMT.` While both involve the concepts of world model and prompt, **Meta-DT differs significantly from CMT in many crucial aspects** (refer to our A2). Also, **other four reviewers** (we refer to iZYy as R1, LRM6 as R2, sbh9 as R4, and BWEr as R5) **are quite affirmative to our 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). We hope that our novelty and distinct contributions have been adequately justified.

- `Empirical comparison to CMT.`We did not include CMT as a baseline since we could not get its source code. Following your advice, **we have added CMT as a new baseline**, and evaluation results again show the superiority of our method. 

- `Evaluation on more environments.` We evaluated our method on **seven environments adopted from three classical benchmarks** in meta-RL. Also, **other four reviewers are quite affirmative to the extensiveness of our empirical evaluation** ("shows the good performance" by R1, "comprehensive experiments on multiple benchmarks" by R2, "strong experimental evaluation over many different tasks" by R4, and "comprehensive experimental results" by R5). Following your advice, **we have conducted new evaluations on Humanoid-Dir and two more Meta-World environments**. Finally, results of **a total of 10 environments from 3 benchmarks** again verify our advantages. We hope that our extended experiments can further verify the generalization of our method.

The corresponding figures of extended experiments can be found in Global Response. Again, we appreciate your time in reviewing our manuscript and your valuable comments. We would like to justify our method to you in more detail, and we are more than delighted to have further discussions to improve our manuscript. If our response has addressed your concerns, we would be grateful if you could re-evaluate our work.