Return-Aware Offline Reinforcement Learning via Multi-modal Sequence Modeling
摘要
Offline Reinforcement Learning (RL) has become an important solution for current intelligent decision-making by learning the optimal policy from pre-existing offline datasets without interacting with the environment. Recent works represented by Decision Transformer (DT) have become a promising method by formulating offline RL as a sequence modeling problem. However, these methods model RL trajectories as one sequence assuming every data belongs to the same modality and ignore the inherent particularity of RL trajectories containing three different modalities of data (state, action, and reward), which can lead to suboptimal performance and make them difficult to efficiently and flexibly cope with diverse types of environments, such as delayed rewards and stochastic dynamics. To explore cross-modal interactions and mitigate susceptibility to environments, we propose a multi-modal transformer, named Return-Aware Transformer Transducer (RATT). Our method comprises three steps. First, we reformulate the principle of DT through the new multi-modal sequence modeling. Second, we design a multi-modal transformer architecture and introduce the return detector to dynamically adapt the return sequence. Third, we train a RATT based on the return detector to generate actions. Finally, during inference, RATT generates actions conditioned on the adapted return sequence. Experiments on various offline RL benchmarks show that RATT effectively learns policies satisfying different environmental characteristics and outperforms baseline methods. We also conduct extensive ablation studies and parameter analysis to highlight the effectiveness of the RATT design.