Q1. Are there any environments or types of tasks where Meta-DT's performance is notably limited? What are the challenges in extending the model to such environments, and how might these be addressed in future work?

A1. Thank you for your insightful questions. We evaluated Meta-DT on seven environments adopted from three classical benchmarks in meta-RL. During rebuttal, we also conducted experiments on new evaluation domains, including Humanoid-Dir and two more Meta-World environments (Figure S2 in Global Response). Results showed the consistent superiority of Meta-DT in these ten environments.

On the other hand, the above environments are relatively homogenous in nature. As stated in the limitations section, our generalist model is trained on relatively lightweight datasets compared to popular large models. The urgent trend that RL practitioners are striving to break through is to deploy on significantly large datasets with more diversity. As mentioned by Reviewer sbh9, it would be interesting to see how Meta-DT generalizes across worlds and tasks with more diversity, like training on $K$-levels of an Atari game and generalizing to the remaining $N-K$ levels of well-suited Atari games.

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. We will include the above insights into our limitations section, and are looking forward to investigating these promising future directions.

Q2. Elaborate on the process and benefits of world model disentanglement in Meta-DT, and how it affects the model's generalization ability.

A2. Existing context-based or prompt-based methods usually infer task representations by feeding subsets of experience into a trajectory encoder. However, in RL regimes, since the offline dataset depends on both the task and the behavior policy, the task information could be entangled with the features of behavior policies, thus producing biased task inference at test time due to the shift in behavior policies. Taking 2D navigation as an example, tasks differ in goals and the behavior policy is going towards the goal for each task. The algorithm can easily distinguish tasks based on state-action distributions rather than reward functions, leading to extrapolating errors when the behavior policy shifts to random exploration during testing.

Hence, we attempt to find a stable way to accurately disentangle task-relevant information from behavior policies. For RL paradigms, the world model completely describes the characteristics of the MDP/task, and keeps invariant to behavior policies or collected datasets. Naturally, it could be a promising alternative for accurately disentangling task beliefs. That is exactly the motivation of our world model disentanglement.

In the meta-learning setting, we assume that the world model shares some common structure across the task distribution. The principle behind is that tasks with similar contexts will behave similarly in the world model. Intuitively, we map tasks into a latent space projected by the world model, and infer task similarity via the world model dynamics. That is, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. Subsequently, we pretrain the world model and then train the meta-policy with a fixed world model. This kind of decoupling assures a more stable process for learning task representations, since the algorithm can solely focus on capturing task similarity via the world model dynamics.

Summary Response

Thank you for your valuable comments, which help us gain more insights into our methodology and motivate promising directions for future work. We are honored to have your recognition of our work, and are also grateful that most reviewers are quite affirmative to our overall contributions, including the novelty and motivation, the extensiveness of empirical evaluation, and the superior performance.

The corresponding figures of extended experiments can be found in the PDF of Global Response. Please let us know if we have addressed your concerns, and we are more than delighted to have further discussions and improve our manuscript.

Q1. It would be nice to see more detailed evaluation in non-mujoco tasks.

A1. Thank you for your advice. We have conducted new evaluations on Humanoid-Dir and two more Meta-World environments of Sweep and Door-Lock. The following table shows the few-shot testing performance of Meta-DT against baselines using Medium datasets (Figure S2 in Global Response). The results again demonstrate the consistent superiority of Meta-DT over all baselines in these newly evaluated domains.

Q2. It would be interesting to see how the technique generalizes across worlds and tasks with more diversity in them, like Atari games.

A2. Thank you for pointing out this insightful direction. As stated in the limitations section, our generalist model is trained on relatively lightweight datasets compared to popular large models. The urgent trend that RL practitioners are striving to break through is to deploy on significantly large datasets with more diversity, unlocking the scaling law with the transformer architecture. Due to limited time during rebuttal, it is hard for us to conduct systematic empirical evaluation on Atari domains. We will include the above insights into our limitations section, and are looking forward to investigating this promising future direction.

Q3. Why do you think the disentanglement works well as an idea?

A3. Thank you for your insightful question. Existing context-based or prompt-based methods usually infer task representations by feeding subsets of experience into a trajectory encoder. The offline dataset depends on both the task and the behavior policy. The task information could be entangled with the features of behavior policies, thus producing biased task inference at test time due to the shift in behavior policies.

Hence, we attempt to find a stable way to accurately disentangle task-relevant information from behavior policies. In the RL regime, the world model completely describes the characteristics of the MDP/task, and keeps invariant to behavior policies or collected datasets. Naturally, it could be a promising alternative for accurately disentangling task beliefs. That is the beginning of motivating the world model disentanglement as our idea.

In the meta-learning setting, we assume that the world model shares some common structure across the task distribution. The principle behind is that tasks with similar contexts will behave similarly in the world model. Intuitively, we map tasks into a latent space projected by the world model, and infer task similarity via the world model dynamics. That is, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. Subsequently, we pretrain the world model and then train the meta-policy with a fixed world model. This kind of decoupling assures a more stable process for learning task representations, since the algorithm can solely focus on capturing task similarity via the world model dynamics.

Summary Response

Thank you for your valuable comments, which help us enhance our experimental evaluation and gain more insightful directions for future work. We are honored to have your recognition of our work, and are also grateful that most reviewers are quite affirmative to our overall contributions, including the novelty and motivation, the extensiveness of empirical evaluation, and the superior performance.

The corresponding figures of extended experiments can be found in the PDF of Global Response. Please let us know if we have addressed your concerns, and we are more than delighted to have further discussions and improve our manuscript.

Q1. The work misses the CMT baseline which is cited but not shown as a baseline.

A1. The reason why we did not include CMT as a baseline is that its code has not been open-sourced. We emailed CMT's authors for requesting the source code, and received the response that the code would be made public at an uncertain time. Following your suggestion, we have included CMT as a new baseline using our own implementation (Figure S1 of Global Response). The results again demonstrate the consistent superiority of Meta-DT over CMT.

Q2. How do the world modeling and prompt tuning in the proposed method differ from CMT?

A2. We do not tune the prompt. Meta-DT selects the trajectory segment using the world model to construct the prompt, without tuning any model parameters associated with the prompt. While both involve the concepts of world model and prompt, Meta-DT differs significantly from CMT in many crucial aspects including:

• The way of using the world model for task inference is different. CMT infers task beliefs at the trajectory level, abstracting a task prompt $z_{\tau}$ from a context trajectory $\tau$ and using that prompt to guide policy generation for the entire episode. In contrast, Meta-DT performs task inference at the transition level, abstracting task representation $z_t$ at each timestep $t$. Hence, Meta-DT can provide more fine-grained guidance to realize high-capacity generalization, and can achieve zero-shot generalization via deriving real-time task inference.

• The prompt design is quite different. CMT abstracts task presentation $z_{\tau}$ from a context trajectory and uses $z_{\tau}$ as the task prompt. In contrast, we directly use a trajectory segment $\tau^*$ as the task prompt, enjoying the power of architecture inductive bias like in Prompt-DT. Moreover, CMT tunes the prompt adaptor layer based on relabeling the offline dataset, while Meta-DT does not involve prompt tuning. Our complimentary prompt design is more lightweight and easy-to-implement. By using the world model to help construct the prompt, the world model (algorithmic perspective) and the complementary prompt (architecture perspective) work as an efficient whole to perform accurate task inference for Meta-DT.

• The pipeline of training the world model and meta-policy is different. CMT simultaneously learns the world model and meta-policy. In contrast, Meta-DT decouples the learning process, which pretrains the world model and then trains the meta-policy with a fixed world model. This kind of decoupling assures a more stable process for learning task representations, since the algorithm can solely focus on capturing task similarity via the world model dynamics.

• The algorithm paradigm is different. Our method is based on decision transformer, a return-conditioned supervised learning paradigm that has achieved promising results on offline RL. It feeds the (return-to-go, state, action) tuple sequence to the causal transformer to predict the action element only. In contrast, CMT uses another paradigm that feeds the (state, action, reward) tuple sequence to the causal transformer to perform autoregressive prediction on all elements.

Q3. Evaluation on Humanoid-Dir and more Meta-World environments.

A3. Following your suggestion, we have conducted new evaluation on Humanoid-Dir and two Meta-World environments of Sweep and Door-Lock. The following table shows the few-shot testing performance of Meta-DT against extended baselines using Medium datasets (Figure S2 in Global Response). The results again demonstrate the consistent superiority of Meta-DT over all baselines in these domains.

Q4. The run-time during training and inference is unclear for fare comparison with baselines.

A4. Please refer to our A1 to Reviewer iZYy for the tables of model parameters, training time, and inference time for one episode. Compared to DT-based baselines, our methods incurrs about 15% increase in model parameters, and about 10% increase in training and inference time. Considering the significant performance gain of our method, this lightweight cost is likely acceptable. Moreover, Meta-DT's computation cost is even lower than other baselines.

Q5. What is the sensitivity with the prompt length $k$?

A5. We have conducted new experiments to investigate the influence of prompt length $k$ on Meta-DT’s performance (Figure S3 in Global Response). Generally, the performance of Meta-DT is not sensitive to the pre-defined value of prompt length.

Q6. The limitations section is inadequate.

A6. We have discussed three threads of our limitations. During rebuttal, the valuable comments from all reviewers also give us new insights into our method, and we will have a more thorough discussion on limitations based on these comments.

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.

Q1. Some intuitive explanation on the extrapolation ability are also desirable.

A1. Thank you for your insightful comments. Meta-DT's generalization comes from the extrapolation ability of the context-aware world model $W(r,s'|s,a; z^i)$. The world model completely describes the characteristics of the MDP/task. In the meta-learning setting, we assume that the world model shares some common structure across the task distribution, with a latent representation $z^i$ to approximate the unknown context of task $i$. The principle behind is that tasks with similar contexts will behave similarly in the world model. Intuitively, we map tasks into a latent space projected by the world model, and infer task similarity via the world model dynamics. That is, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. For discrete tasks like Cheetah-Dir, the world model also shares some common structure across the task distribution, e.g., the kinematics principle or locomotion skills. Hence, the extrapolation of the world model works for both continuous and discrete tasks. We will enhance the intuitive explanation of our method in the revised paper.

Q2. It may be helpful to provide the experiments with different numbers of training tasks.

A2. Following your advice, we have conducted new experiments to evaluate Meta-DT with a different number of training tasks. The following table shows the few-shot test performance of Meta-DT using Medium datasets (Figure S4 in Global Response). Increasing the number of training tasks can usually improve the generalization ability to test tasks, especially in harder tasks like Ant-Dir.

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. We will include this new experiment and the interesting findings in our revised paper.

Q3. It is also interesting to see the extrapolation ability limit of meta-DT.

A3. Following your suggestion, we have conducted new experiments on Ant-Dir to evaluate Meta-DT's generalization ability to out-of-distribution (OOD) tasks. The following table shows the few-shot performance of Meta-DT with OOD test tasks (Figure S5 in Global Response). In this case, we sample training tasks with a goal direction of $U[0, 1.5\pi]$, and then test on tasks of $U[1.5\pi, 2\pi]$.

Obviously, Meta-DT can still obtain better performance than baselines on OOD test tasks, which again verifies our superiority. As stated in our A1 to Q1, we extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model, which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. For tasks like Ant-Dir, the world model shares some common structure across the task distribution (even for OOD tasks), e.g., the kinematics principle or locomotion skills. Hence, the extrapolation of the world model also works for OOD test tasks in this case. We will include this interesting experiment in the revised paper.

Summary Response

Thank you for your valuable suggestions, which help us gain more critical insights into the extrapolation ability of Meta-DT, and further enhance our experimental evaluation. The corresponding figures of extended experiments can be found in the PDF of Global Response. We are honored to have your recognition of our method, and are also grateful that most reviewers are quite affirmative to our overall contributions, including the novelty and motivation, the extensiveness of empirical evaluation, and the superior performance.

Please let us know if we have addressed your concerns. We are more than delighted to have further discussions and improve our manuscript.

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.

Q1. Explain more about the training hours and parameters for this method.

A1. The following tables show the number of model parameters, the training time, and inference time for one episode. Compared to DT-based baselines, our method introduces a lightweight world model that consists of about 15% parameters of the DT backbone. Thanks to this lightweight design, Meta-DT incurrs about 10% runtime increase during training and inference. Considering the significant performance gain of our method, this lightweight cost is likely acceptable. Moreover, Meta-DT's computation cost is even lower than other baselines.

Q2. Explain the experiments why the method does not outperform the baselines in the expert demonstration.

A2. In experiments with Expert datasets as shown in Table 4 and Figure 10, Meta-DT outperforms three baselines (Generalized DT, CORRO, and CSRO) to a large extent in all environments.

Compared to Prompt-DT, Meta-DT obtains significantly better performance in Point-Robot (-6.90 vs. -7.99), Cheetah-Vel (-52.42 vs. -133.78), and Ant-Dir (961.27 vs. 678.07), and obtains extremely close performance in Cheetah-Dir (874.91 vs. 960.32), Hopper-Param (383.51 vs. 393.79), and Walker-Param (437.79 vs. 449.15). Taken together, it can be empirically verified that Meta-DT still outperforms Prompt-DT when evaluated on Expert datasets.

With Expert datasets, Prompt-DT can access high-quality expert demonstrations as the task prompt to guide policy generation. Within the expert data, the behavior policy exactly equals the optimal policy and totally aligns with task characteristics, thus it can capture task-relevant information very well. Hence, Prompt-DT can achieve the same level of performance as Meta-DT in several environments. When faced with inferior datasets, the task information could be entangled with the features of behavior policies, producing biased task inference at test time. Hence, the performance of Prompt-DT might drop a lot, while Meta-DT can still obtain satisfactory performance due to accurately capturing task information via world model disentanglement.

Q3. This method is not as light as the previous work like Prompt-DT.

A3. We agree with you on this point, as our method introduces a pretrained context-aware world model. In contrast, Prompt-DT uses the same network architecture as DT, with some expert demonstrations prepended to the DT's input.

Our motivation of introducing the additional world model is to disentangle task-relevant information from behavior policies, thus more accurately inferring task beliefs. With this effective disentanglement, our method is verified to be robust to the dataset quality and is more practical with fewer prerequisites in real-world scenarios. This is what Prompt-DT lacks, since it is sensitive to the quality of prompt demonstrations and can achieve satisfactory performance only when expert demonstrations are available at test time.

Summary Response

Thank you for your valuable review comments, which help us gain more insights on comparison to existing works like Prompt-DT, and further enhance our experimental illustration. We are honored to have your recognition on our method, especially on its novelty ("a novel Meta-DT method") and superior performance ("shows the good performance... outperform the baselines").

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

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

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