Title: Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement
Venue: NeurIPS 2024
Accepted: yes
Keywords: decision transformer

offline meta reinforcement learning

world model
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

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.

1 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

self-supervised pretraining techniques [18, 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 to memorize the features of behavior policies and produce biased task inference at test time due to the change of behavior policies [24, 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 of the underlying environment [26]. The environment dynamics, also called world model [27, 28], is intrinsically invariant to behavior policies or collected datasets, thus presenting a 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 [31, 32].

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.

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

3 Preliminaries

3.1 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 as max π 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 as J ( π meta ) = E M ∼ P ( M ) [ J M ( π meta )] .

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 [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 = ( · · · , ˆ R t , s t , a t , · · · ) , where ˆ R t 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, ˆ R t = ∑ T i = t r i , and during testing, ˆ R t = 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 ˆ R t , 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.

Multi-Task OfflineDataset

ut

~Di

a

米

α t-K+1

at-1

a

t

Meta Decision Transformer

S

a

r

t-h

S

a

r

:t-1

S

t

Context-Aware

Context Encoder

Z

RS

a

Z

a

Z

R

S

World Model

E(2μ)

t-K+1

t-1

t

j+k

R

*

S

*

a

*

Complementary Prompt

Reward Decoder

TransitionDecoder

R(rt|st,at; 2t

Tp(st+1lst,at; 2t

j+k

HistoryTrajectories

>

C

j

R

S

a

0

R

S

a

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

4.1 Context-Aware World Model

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 as

Since 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

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 ˆ r t +1 = R ϕ ( s t , a t ; z i t ) and next state ˆ s 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 as

4.2 Complementary Prompt

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

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.

4.3 Meta-DT Architecture

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 as

The input sequence ( τ ∗ i , τ + i ) corresponds to 3 k + 4 K 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.

5 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 MetaWorld 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.

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 warm-start data to infer the task representation. We first compare MetaDT 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

CORRO

Prompt-DT

GeneralizedDT

CSRO

Meta-DT (Ours)

Point-Robot

Cheetah-Vel

Cheetah-Dir

-10

-100

600

Returns

15

wm

Returns

-150

Returns

400

-20

-200

200

-25

-250

0

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

Ant-Dir

Hopper-Param

Walker-Param

450

400

350

400

MMMM

Returns

300

Returns

Returns

300

350

N

250

250

200

200

MN

200

150

0

10

2030

40

50

0

2

4

6

8

10

0

2

4

6

8

10

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

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

CORRO

Prompt-DT

Generalized DT

CSRO

Meta-DT (Ours)

Point-Robot

Cheetah-Vel

Cheetah-Dir

600

-10

-100

-15

400

Returns

-20

Returns

-150

Returns

200

-25

-200

-30

-250

0

-35

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

Ant-Dir

Hopper-Param

Walker-Param

450

400

350

400

Returns

300

Returns

300

Returns

350

200

300

250

vW

250

100

200

0

200

150

0

10

20

30

40

50

0

2

4

6

8

10

0

2

4

6

8

10

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

5.3 Ablation Study

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.

5.4 Robustness to the Quality of Offline Datasets

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. Figs. 5-6 present test return curves of

DT

w/ocontext

w/oprompt

w/ocomplementary

Meta-DT

Point-Robot

Cheetah-Dir

Ant-Dir

-10

600

400

-12

Returns

Returns

S

400

eturns

S

300

-14

200

-16

200

-18

100

0

1

2

3

4

5

0

0.0

0.2

0.4

0.6

0.8

1.0

0

1

2

3

4

5

Timesteps (1e4)

Timesteps (1e5)

Timesteps (1e5)

CORRO

Prompt-DT

Generalized DT

CSRO

Meta-DT (Ours)

Point-Robot

Cheetah-Vel

Cheetah-Dir

-10

-75

600

-100

Returns

-15

Returns

-125

Returns

400

-20

Mh

-150

200

-25

-175

-200

0

0

10

20

30

0

5

10

15

20

0

10

20

30

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

1000

Ant-Dir

375

Hopper-Param

500

Walker-Param

800

350

MMM

450

Returns

600

Returns

400

350

400

300

300

200

275

250

0

0

10

20

30

40

50

250

0

2

4

6

8

10

200

0

2

4

6

8

10

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

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.

CORRO

Prompt-DT

Generalized DT

CSRO

Meta-DT (Ours)

Point-Robot

Cheetah-Vel

Cheetah-Dir

-50

1000

-10

Returns

Returns

-100

Returns

800

15

600

-150

-20

MWW

400

m

-200

200

-25

0

10

20

30

0

2

4

6

8

10

0

10

20

30

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

Ant-Dir

Hopper-Param

Walker-Param

1000

400

450

800

400

Returns

Returns

350

Returns

350

600

300

300

400

250

200

250

200

0

10

20

40

50

0

2

4

6

8

10

0

2

4

9

8

10

Timesteps (1e4)

Timesteps (1e4)

Timesteps (1e4)

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

References

Jacob Devlin, Ming-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.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. ViViT: A video vision transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 6836-6846, 2021.

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.

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.

Thomas Schmied, Markus Hofmarcher, Fabian Paischer, Razvan Pascanu, and Sepp Hochreiter. Learning to modulate pre-trained models in rl. In Advances in Neural Information Processing Systems , 2023.

Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Proceedings of Conference on Robot Learning , pages 785-799, 2023.

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.

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.

David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? In Advances in Neural Information Processing Systems , volume 35, pages 1542-1553, 2022.

Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. In Advances in Neural Information Processing Systems , 2023.

Anirudhan Badrinath, Yannis Flet-Berliac, Allen Nie, and Emma Brunskill. Waypoint transformer: Reinforcement learning via supervised learning with intermediate targets. In Advances in Neural Information Processing Systems , volume 36, 2023.

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.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , pages 5998-6008, 2017.

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 16x16 words: Transformers for image recognition at scale. In Proceedings of International Conference on Learning Representations , 2021.

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.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In Proceedings of International Conference on Machine Learning , pages 8821-8831, 2021.

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.

Runji Lin, Ye Li, Xidong Feng, Zhaowei Zhang, Xian Hong Wu Fung, Haifeng Zhang, Jun Wang, Yali Du, and Yaodong Yang. Contextual transformer for offline meta reinforcement learning. In NeurIPS 2022 Foundation Models for Decision Making Workshop , 2022.

Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt-tuning decision transformer with preference ranking. arXiv preprint arXiv:2305.09648 , 2023.

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.

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 on Machine Learning , volume 162, pages 24631-24645, 2022.

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.

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.

Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcement learning via contrastive learning. In Proceedings of International Conference on Machine Learning , pages 25747-25759, 2022.

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 meta-reinforcement learning. In Advances in Neural Information Processing Systems , volume 36, 2023.

Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction . The MIT Press, second edition, 2018.

David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122 , 2018.

Tim Brooks, Bill Peebles, Connor Homes, Will DePue, Yufei Guo, et al. Video generation models as world simulators. 2024.

Arunkumar Byravan, Jost Tobias Springenberg, Abbas Abdolmaleki, 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.

Samuel Kessler, Mateusz Ostaszewski, Michał PawełBortkiewicz, Mateusz ˙ Zarski, Maciej Wolczyk, Jack Parker-Holder, Stephen J. Roberts, and Piotr Mi´ s. The effectiveness of world models for continual reinforcement learning. In Proceedings of Conference on Lifelong Learning Agents , volume 232, pages 184-204, 2023.

Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. In Advances in Neural Information Processing Systems , volume 34, pages 28954-28967, 2021.

Zohar Rimon, Tom Jurgenson, Orr Krupnik, Gilad Adler, and Aviv Tamar. MAMBA: an effective world model approach for meta-reinforcement learning. In Proceedings of International Conference on Learning Representations , 2024.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Alexander Nikulin, Vladislav Kurenkov, Denis Tarasov, and Sergey Kolesnikov. Antiexploration by random network distillation. In Proceedings of International Conference on Machine Learning , 2023.

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.

Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang. MetaDiffuser: Diffusion model as conditional planner for offline meta-rl. In Proceedings of International Conference on Machine Learning , 2023.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of International conference on machine learning , pages 1126-1135, 2017.

Zhongwen Xu, Hado P van Hasselt, and David Silver. Meta-gradient reinforcement learning. In Advances in Neural Information Processing Systems , volume 31, 2018.

Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline metareinforcement learning with advantage weighting. In Proceedings of International Conference on Machine Learning , pages 7780-7791, 2021.

Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In Proceedings of International Conference on Machine Learning , pages 5331-5340, 2019.

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):13198-13236, 2021.

Lanqing Li, Rui Yang, and Dijun Luo. FOCAL: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization. In Proceedings of International Conference on Learning Representations , 2021.

Lanqing Li, Hai Zhang, Xinyu Zhang, Shatong Zhu, Junqiao Zhao, and Pheng-Ann Heng. Towards an information theoretic framework of context-based offline meta-reinforcement learning. arXiv preprint arXiv:2402.02429 , 2024.

Zhi Wang, Chunlin Chen, and Daoyi Dong. Lifelong incremental reinforcement learning with online Bayesian inference. IEEE Transactions on Neural Networks and Learning Systems , 33(8):4003-4016, 2022.

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.

William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 4195-4205, 2023.

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.

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):233-248, 2023.

Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-maron, 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.

Zhe Wang, Haozhu Wang, and Yanjun (Jane) Qi. T3GDT: Three-tier tokens to guide decision transformer for offline meta reinforcement learning. In NeurIPS 2023 Workshop on Robot Learning , 2023.

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.

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.

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.

Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys , 55(9):1-35, 2023.

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.

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-World: A benchmark and evaluation for multi-task and meta reinforcement learning. In Proceedings of Conference on Robot Learning , pages 1094-1100, 2020.

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.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

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.

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.

Appendix A. Algorithm Pesudocodes

Based on the implementations in Sec. 4, this section gives the brief procedures of our method. First, Algorithm 1 presents the pretraining of the context-aware world model. Then, Algorithm 2 shows the pipeline of training Meta-DT, where the sub-procedure of generating the complementary prompt is given in Algorithm 3. Finally, Algorithm 4 and Algorithm 5 show the few-shot and zero-shot evaluations on test tasks, respectively.

Algorithm 2: Training Meta Decision Transformer

Algorithm 3: Complementary Prompt Generation ( GetPrompt )

Pretrained context encoder E ψ and decoders R ϕ , T φ

1 Sample a trajectory ( s 0 , a 0 , r 0 , ..., s t , a t , r t , ... ) from D demo i

2 Get the segment with the largest prediction error on the pretrained world model as

3 Obtain the k -step prompt τ ∗ 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 )

Algorithm 4: Meta-DT Few-Shot Evaluation on Test Tasks

Algorithm 5: Meta-DT Zero-Shot Evaluation on Test Tasks

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 [39, 24], we adopt three classical benchmarks: the 2D navigation [25], the multi-task MuJoCo control [55, 36], and the Meta-World [56]. We evaluate all tested methods on the following environments as

Point-Robot : a problem of a point agent navigating to a 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 is [ -0 . 1 , 0 . 1] 2 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 a goal. Cheetah-Dir/Ant-Dir is to control a 2D cheetah/3D quadruped ant robot to move in a specific direction, and the reward function is the cosine product of the agent's velocity and the goal direction. The goal velocity and direction are uniformly sampled from the distribution U [0 . 075 , 3 . 0] for Cheetah-Vel and U [0 , 2 π ] for Ant-Dir. In Cheetah-Dir, the goal directions are limited to forward and backward. The maximal episode step is set to 200.

Hopper-Param , Walker-Param : multi-task MuJoCo benchmarks where tasks differ in state transition dynamics. The two benchmarks control a one-legged hopper and a two-legged walker robot to run as fast as possible. The reward function is proportional to the running velocity in the positive x -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 Meta-World. 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 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, we sample 45 tasks for training and another 5 held-out tasks for testing. For Meta-World environments, we sample 15 tasks for training and 5 held-out tasks for testing.

B.2. The Details of Datasets

We employ the soft actor-critic (SAC) [57] algorithm to train a policy independently for each task. Table 5 shows the detailed hyperparameters for training the SAC agents in each environment. Figs. 7-8 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 a medium policy that achieves a one-third to one-second score of an expert policy. All trajectories are generated by this medium policy.

Expert : using an 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.

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.

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

Appendix C. The Details of Baselines

This section gives the details of the representative five baselines, including three DT-based approaches, one temporal difference-based approach, and one diffusion-based approach. These baselines are thoughtfully selected to encompass the three distinctive categories of offline meta-RL methods. Also, since our method Meta-DT belongs to the DT-based category, we adopt more approaches from this kind as our baselines. The baselines are introduced as follows:

Prompt-DT [21], is a DT-based 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 can assess a handful of expert demonstrations to construct the prompt.

Generalized DT [22], is a DT-based method that formulates multi-task learning as a 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 statemarginal matching and imitation learning as two generic HIM problems to evaluate the proposed Categorical DT and Bi-directional DT.

Point-Robot(Train Tasks)

Half-Cheetah-Vel (Train Tasks)

Half-Cheetah-Dir(Train Tasks)

0

-5

-100

1250

Returns

-10

Returns

-200

Returns

1000

-300

750

-20

-400

500

-25

-500

250

-30

-600

0

0.0

0.5

1.0

2.0

0

1

2

3

4

5

0

1

2

3

4

Timesteps (1e3)

Timesteps (1e5)

Timesteps (1e5)

Point-Robot(Test Tasks)

Half-Cheetah-Vel(Test Tasks)

Half-Cheetah-Dir(Test Tasks)

0F

-5

1250

-100

1000

Returns

-10

Returns

-200

Returns

750

-15

500

-300

250

-20

-400

0

0.0

0.5

1.0

1.5

2.0

0

1

2

3

4

5

0

1

2

3

4

5

Timesteps (1e3)

Timesteps (1e5)

Timesteps (1e5)

Ant-Dir (Train Tasks)

Hopper-Rand-Params(Train Tasks)

Walker-Rand-Params (Train Tasks)

1200

1000

800

1000

Returns

500

750

H

Returns

600

400

Returns

800

600

250

200

400

0

OW

0

200

-250

0

1

2.3

4

5

0

1

2.34

5

0

1

2.3

4

5

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

Ant-Dir(Test Tasks)

Hopper-Rand-Params(Test Tasks)

Walker-Rand-Params(Test Tasks)

700

1000

1000

600

M

800

500

800

Returns

Returns

Returns

600

400

600

300

400

400

200

200

100

200

0

1

4

5

0

1

23

4

5

0

1

4

5

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

Reach(Train Tasks)

Sweep (Train Tasks)

DoorLock(Train Tasks)

5000

4000

4000

3000

3000

Returns

3000

Returns

Returns

2000

2000

2000

1000

1000

1000

0

0.0

0.5

1.0

1.5

2.0

0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

Reach (Test Tasks)

Sweep(Test Tasks)

DoorLock (Test Tasks)

5000

1000

4000

4000

800

3000

Returns

3000

Returns

600

Returns

2000

400

2000

1000

200

1000

0

0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

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

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 a GRU network with ReLU as the activation function. The GRU cell encodes the agent's h -step history ( s t -h , a t -h , r t -h , ..., s t -1 , a t -1 , r t -1 , s t ) into a 128-dimensional vector, and the MLP transforms this vector into a 16-dimensional embedding, i.e., the task representation z . The reward decoder is an MLP that takes as input the ( s, a, s ′ , z ) tuple and predicts the reward r through two 128-dimensional hidden layers. Similarly, the state transition decoder is an MLP that takes as input the ( s, a, r, z ) 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 [8] (https://github.com/kzl/decision-transformer). We follow most of the hyperparameters as they did. Specifically, task representations z , returns ˆ R , states s , and actions a are fed into modality-specific linear embeddings, and the same positional episodic timestep encoding is added to 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 6 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 on one environment (PointRobot, Cheetah-Vel, Cheetah-Dir, Ant-Dir, Hopper-Param, Walker-Param, Reach) with one dataset (Medium, Expert, Mixed). Some common hyperparameters across all report units are set as: optimizer Adam, weight decay 1e-4, linear warmup steps for learning rate decay 10000, gradient norm clip 0.25, dropout 0.1, and batch size 128. Table 7 presents the detailed hyperparameters of Meta-DT trained on the Point-Robot and MuJoCo domains with the Medium, Expert, and Mixed datasets. Table 8 presents the detailed hyperparameters of Meta-DT trained on Meta-World environments with the Medium datasets.

Compute. We train our models on one Nvidia RTX4080 GPU with the Intel Core i9-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.

Appendix E. Hyperparameter Analysis

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.

Fig. 9 and Table 9 present the testing curves and converged performance of Meta-DT with different values of context horizon h on representative environments. Generally, the performance of Meta-DT

is not sensitive to the pre-defined value of the context horizon. In simple environments like PointRobot, 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 a 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 a lightweight network design while not sacrificing performance.

horizon-4

horizon-6

horizon-8

Point-Robot

Cheetah-Dir

Ant-Dir

700

450

-10

600

400

Returns

-12

eturns

500

eturns

350

400

300

R

300

-14

200

250

-160.0

1000.0

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

200

0

1

2

3

4

5

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

Prompt length k . Analogously, we also investigate the influence of prompt length k on Meta-DT's performance. Fig. 10 and Table 10 present the testing curves and converged performance of Meta-DT with different values of prompt length k on representative environments. Generally, the performance of Meta-DT is not sensitive to the pre-defined value of prompt length.

prompt-3

prompt-5

prompt-7

Point-Robot

700

Cheetah-Dir

450

Ant-Dir

-10

3333

600

400

Returns

-12

Returns

500

ns

350

400

R

300

R

300

-14

200

250

-160.0

1000.0

200

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

0

1

2

3

4

5

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

The number of training tasks . We evaluate Meta-DT with a different number of training tasks. Fig. 11 and Table 11 present the testing curves and converged performance of Meta-DT with different numbers of training tasks on 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 '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.

n_task-15

n_task-30

n_task-45

Point-Robot

Cheetah-Vel

Ant-Dir

450

-10

-100

400

Returns

-12

Returns

eturns

350

-150

300

-14

-200

250

-160.0

-2500.0

200

0.2

0.4

0.6

0.8

1.0

0.2

0.4

0.6

0.8

1.0

0

1

2

3

4

5

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

Appendix F. Experimental Results on Meta-World

Here, we present experimental results on the robotic manipulation benchmark Meta-World [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

GeneralizedDT

CSRO

Meta-DT (Ours)

Reach

Sweep

Door-Lock

3500

800

2500

3000

600

2500

Returns

Returns

eturns

2000

2000

400

N

E

1500

200

MMM

1500

1000

M

WW

1000

0.0

0.5

1.0

1.5

0.0

0.5

1.0

1.5

500

0.0

0.2

0.4

0.6

0.8

Timesteps (1e5)

Timesteps (1e5)

Timesteps (1e5)

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 in a 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 on Ant-Dir to evaluate Meta-DT's generalization ability to OOD tasks. In this case, we sample training tasks with a goal direction of U [0 , 1 . 5 π ] and then test on tasks of [1 . 5 π, 2 π ] . Fig. 13 and Table 13 show the few-shot return curves and converged performance on OOD test tasks using Medium datasets from Ant-Dir.

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

Ant-Dir

500

CORRO

Prompt-DT

400

Generalized DT

CSRO

Meta-DT (Ours)

eturns

S

300

R

200

100

0

1

2

3

4

5

Timesteps (1e5)

NeurIPS Paper Checklist

1. 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 1: Introduction.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Please refer to 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Please refer to Appendix A for our algorithm pseudocodes, Appendix B for the dataset construction, and Appendix D 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have provided the source code with sufficient instructions at github: https://github.com/NJU-RL/Meta-DT.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Please refer to Section 5: Experiments.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Please refer to Appendix D.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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.

11. Safeguards

Question: 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)?

Answer: [NA]

12. Licenses for existing assets

Question: 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?

Answer: [NA]

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

14. Crowdsourcing 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 as 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nA 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.\n\nSTRENGTHS:\nGeneralization 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.\n\nWEAKNESSES:\n1. 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?\n\nQUESTIONS:\nMentioned in the weakness section\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\nNo 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.\n\nQUESTIONS:\nWhy 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?\n\nLIMITATIONS:\nThe limitations and future work are well addressed", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

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

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

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:
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:
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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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$].\n\nQUESTIONS:\nPlease also consider responding to the Weaknesses.\n\nLIMITATIONS:\nThe authors have discussed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

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

STRENGTHS:
- This paper proposes a novel Meta-DT that disentangles the task information from the trajectories as the prompt in Meta RL tasks.
- 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.
- This paper shows the sub-performance compared with prompt DT when the demonstration set is expert.

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:
See weaknesses and questions.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-15