Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement

Summary

A new framework is proposed for offline meta-reinforcement learning (meta-RL) that leverages transformers and world model disentanglement to enhance task generalization without the need for expert demonstrations or domain knowledge. The approach, called Meta Decision Transformer (Meta-DT), utilizes a context-aware world model to encode task-relevant information which guides the transformer in generating task-oriented sequences. Experimental results show that Meta-DT achieves superior few-shot and zero-shot generalization across different benchmarks, indicating its practicality and effectiveness in RL applications.

Strengths

Generalization without Experts: Meta-DT effectively generalizes across unseen tasks without requiring expert demonstrations or domain knowledge, which is a significant advancement in the field of meta-reinforcement learning. World 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. Efficient 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. Empirical Validation: The paper includes comprehensive experimental results showing that Meta-DT outperforms existing baselines in few-shot and zero-shot generalization tasks across multiple benchmarks, demonstrating its practical effectiveness.

Weaknesses

1. Are there any environments or types of tasks where Meta-DT's performance is notably limited? What are the challenges in extending the model to such environments, and how might these be addressed in future work? 2. Can you elaborate on the process and benefits of world model disentanglement in Meta-DT? How does this process affect the model's ability to generalize across different tasks, especially in dynamically changing environments?

Questions

Mentioned in the weakness section

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

The paper introduces a new architecture for Meta-RL based on Decision Transformers. The new architecture uses a world model responsible for efficiently encoding task information from demonstrations. The model also introduces a prompt encoder which acts as a boosting mechanism to the context encoder to enhance the models task adaptation abilities.

Strengths

- Strong experimental evaluation including over many different tasks, settings,f ablations, and comparison against several other techniques - Strong results in this evaluation, with good demonstration that the introduced components make a difference. - Simple design with good results - Good figures and explanations, with std error ranges -- which can be critical for RL

Weaknesses

No major flaws with the paper. The evaluation covers environments which are relatively homogenous in nature. Performance in MetaWorld is much less differentiating compared to MuJoCo. It would be nice to see more detailed evaluation in non-mujoco tasks. It would be interesting to see how the technique generalizes across worlds and tasks with more diversity in them, mujoco worlds are extremely minimal with little variation. Something like training on K-levels of an Atari game and generalizing to the remaining N-K levels of well suited Atari games. It is entirely possible the research in this simply isn't ready to tackle hard cases such as this though. Or something like the worlds seen in the Muesli paper.

Questions

Why do you thin the disentanglement works well as an idea? I understand the different perspectives of adaptation the prompt vs context encoder aim to target, but could you not form the context in manner to do this, for example running it twice with different subsets of experience? Is it the way things are fed in to the prompt vs context encoder (segmenting) and the lack of constraints on the prompt encoders representation?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations and future work are well addressed

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$?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

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.

Summary

This paper proposes a novel meta-RL framework called Meta Decision Transformer (Meta-DT). It leverages robust task representation learning via world model disentanglement to achieve task generalization in offline meta-RL. Firstly, it pretrains a context-ware world model to capture the task-relevant information from the offline dataset. Then, it guides the sequence generation of decision transformer using the task representation and the self-guided prompt from past trajectories. The authors conduct extensive experiments on multiple benchmarks. Meta-DT demonstrates better few- and zero-shot generalization ability than other baselines.

Strengths

1. This paper has a clear motivation in disentangling the task specific information. It is quite novel and inspiring to make use of the prompt to provide task-specific context and exploit architecture inductive bias. 2. 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. 3. The paper is well-organized and easy to understand. The authors also provide detailed pseudo-codes and hyper-parameters in the appendix for reproduction.

Weaknesses

1. From my perspective, the generalization ability of meta-DT depends heavily on the extrapolation ability of the context encoder. Although it is intuitive for some continuous tasks like Cheetah-Vel, this extrapolation tends to be quite hard for some other discrete tasks like Cheetah-Dir. Despite the good experiment performance, some intuitive explanation and theoretical justification are also desirable. 2. To evaluate the generalization ability more comprehensively, it may be helpful to provide the experiments with different number of training tasks. It is also interesting to see the extrapolation ability limit of meta-DT. For example, what if the training tasks of Ant-Dir are sampled from U[0,1.5$\pi$] with hold-out tasks are from U[1.5$\pi$, 2$\pi$].

Questions

Please also consider responding to the Weaknesses.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed the limitations.

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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses and questions.

Thank you!

Thank you for your time in helping us improve our work. We are more than delighted to see that we could address your all concerns. We sincerely appreciate you raising the score on our work!