Meta-DT Transforms Reinforcement Learning With Superior Task Generalization

By Dr Silpaja Chandrasekar, PhDReviewed by Joel ScanlonOct 22 2024

Meta-DT leverages cutting-edge transformers and self-guided prompts to push the boundaries of generalization in reinforcement learning, outperforming strong baselines without relying on expert demonstrations.

Study: Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement. Image Credit: metamorworks / Shutterstock

*Important notice: arXiv publishes preliminary scientific reports that are not peer-reviewed and, therefore, should not be regarded as definitive, used to guide development decisions, or treated as established information in the field of artificial intelligence research.

In an article recently submitted to the arXiv preprint* server, researchers at the Chinese Academy of Sciences and Nanjing introduced the meta-decision transformer (Meta-DT), addressing the generalization challenges in reinforcement learning (RL). The framework utilized a sequential modeling ability with a context-aware world model for task representation and guided sequence generation through a causal transformer.

Meta-DT leveraged history trajectories from the meta-policy as self-guided prompts, enhancing task-specific information encoding. These complementary prompts were selected based on the largest prediction error in the pre-trained world model, ensuring they provided additional task-specific information. Experimental results demonstrated that Meta-DT achieved superior few and zero-shot generalizations on multi-joint dynamics with contact (MuJoCo) and meta-world benchmarks compared to strong baselines, requiring no expert demonstrations at test time.

Related Work

Past work in offline Meta-RL focused on learning optimal policies from pre-collected datasets without online interactions. Various methods, including optimization-based and context-based meta-learning, aim to generalize new tasks by utilizing offline task distributions. However, many relied on hand-designed heuristics and faced challenges from the "deadly triad" of RL (instability due to function approximation, off-policy learning, and bootstrapping).

The introduction of the transformer architecture for sequence modeling, as seen in models like decision transformers, has paved the way for new approaches in RL, leveraging autoregressive action outputs and enabling effective multi-task learning. However, models like Prompt-DT still relied on expert demonstrations at test time, limiting their practicality for real-world applications.

Meta-DT Framework

The Meta-DT represents a novel offline Meta-RL framework that integrates advances in conditional sequence modeling with robust task representation learning. This framework aims for efficient generalization across unseen tasks. The method incorporates a context-aware world model designed to accurately encode task-relevant information, which is crucial for extrapolating knowledge. Unlike supervised or unsupervised learners who analyze fixed datasets, Meta-DT addresses the unique challenges of RL, which involves learning optimal policies from dynamic interactions with the environment.

The context-aware world model learns robust task representations that adapt to varying transition dynamics to disentangle task-relevant information from behavior policies. This model approximates the reward and state transition functions across tasks through a shared latent representation. The context encoder captures recent experiences to infer task representations, while the decoder predicts rewards and next states. By accurately encoding task dynamics, the model can extrapolate meta-level knowledge and enhance task generalization capabilities.

In addition to the world model, the framework incorporates a complementary self-guided prompt, which utilizes historical trajectories generated by the meta-policy during evaluation. This prompt is a high-quality demonstration, enabling the model to generalize more effectively without relying on expert knowledge. The segments used for the prompt are selected based on their prediction error in the world model, ensuring that they provide task-relevant information that the world model may not capture fully.

The Meta-DT architecture builds on the DT model and integrates the context-aware world model to solve offline Meta-RL through conditional sequence modeling. The model uses the context encoder to derive contextual information during training, combining a K-step complementary prompt with a recent K-step history.

The input sequence consists of tokens corresponding to these prompts and histories, allowing the model to output actions based on state tokens autoregressively. During testing, Meta-DT can interact with the environment in a few-shot setting to construct a self-guided prompt. In contrast, it evaluates performance directly without the prompt component in zero-shot settings.

Evaluating Meta-DT Generalization Performance

The experiments sought to evaluate the Meta-DT framework's generalization capabilities across various benchmark domains and dataset types. The key questions addressed included whether Meta-DT could consistently outperform strong baselines in few-shot and zero-shot generalization to unseen tasks, how different components like the context-aware world model and complementary prompt design influenced performance, and the robustness of Meta-DT to the quality of offline datasets.

The evaluations were conducted in classical meta-RL environments, such as the point-robot navigation environment, multi-task MuJoCo control, and the MetaWorld manipulation platform. Each environment involved a random sampling of tasks divided into training and testing sets, with datasets constructed using three approaches: medium, mixed, and expert.

Meta-DT was compared against four competitive baselines, including the DT-based prompt-DT and generalized DT, as well as temporal difference-based CORRO and causal structure reinforcement optimization (CSRO). Notably, Meta-DT performed exceptionally well in more complex environments like Ant-Dir, demonstrating its strong generalization capacity for challenging tasks. The results indicated that Meta-DT consistently achieved superior data efficiency and outcomes across different environments, particularly in more complex tasks like ant-dir.

The analysis revealed that Meta-DT demonstrated lower variance during learning, which indicated improved training stability and efficiency compared to the baselines. In the few-shot generalization tests, Meta-DT outperformed its counterparts, particularly in environments with varying reward functions and state transition dynamics.

Regarding zero-shot generalization, Meta-DT showed a minimal drop in performance compared to its few-shot counterparts, demonstrating its ability to derive real-time task representations without relying on expert demonstrations. Its performance drop in zero-shot settings was notably smaller (around 5%) compared to the significant drops seen in other baselines (up to 35%). Ablation studies further confirmed the necessity of each component within the Meta-DT architecture, highlighting the critical role of task representation learning via the world model.

Additionally, the framework proved robust across datasets of varying quality, significantly outperforming baselines, especially in complex environments. These findings emphasize Meta-DT's practicality and effectiveness in real-world scenarios with fewer prerequisites than other methods.

Conclusion

To sum up, the paper addressed the offline Meta-RL challenge by utilizing scalable transformers for improved few and zero-shot generalization. The method demonstrated practical applicability without needing expert demonstrations, although it required a two-phase process for training.

Future work aims to develop a unified framework for concurrent task representation and meta-policy learning, which could enhance the model's efficiency. Additionally, the team is exploring larger datasets and self-supervised learning techniques, aligning with current trends in reinforcement learning practices.

*Important notice: arXiv publishes preliminary scientific reports that are not peer-reviewed and, therefore, should not be regarded as definitive, used to guide development decisions, or treated as established information in the field of artificial intelligence research.