Adaptive Data Augmentation Transformer for Offline Meta-Reinforcement Learning

Offline meta-reinforcement learning (RL) seeks to enable rapid adaptation to new tasks from static datasets. Decision Transformer (DT) leverages the Transformer architecture's ability to model long-range dependencies in sequential data, offering significant promise for offline meta-RL. However, DT's general-ization to unseen tasks remains limited due to distribution shifts and sensitivity to dataset quality. To address the challenges above, we propose Adaptive Data Augmentation Transformer (ADAT), a novel approach that combines adaptive data augmentation with a transformer-based architecture and a world model. ADAT dynamically selects and emphasizes informative trajectories by masking low-advantage experiences, thereby refining the training data distribution. This allows the transformer network to focus on extracting critical temporal patterns, while the integrated world model facilitates counterfactual reasoning. Comprehensive experiments conducted in the Gym MuJoCo environment under few-shot and zero-shot settings demonstrate that ADAT surpasses existing baselines in terms of training convergence rate, asymptotic performance, and generalization capability.