Switch Trajectory Transformer with Distributional Value Approximation for Multi-Task Reinforcement Learning

arXiv 2203.07413 ↗First posted 14 Mar 2022cs.LG cs.AIcited by 11

Not on the map. Why?

Abstract (arXiv)

We propose SwitchTT, a multi-task extension to Trajectory Transformer but enhanced with two striking features: (i) exploiting a sparsely activated model to reduce computation cost in multi-task offline model learning and (ii) adopting a distributional trajectory value estimator that improves policy performance, especially in sparse reward settings. These two enhancements make SwitchTT suitable for solving multi-task offline reinforcement learning problems, where model capacity is critical for absorbing the vast quantities of knowledge available in the multi-task dataset. More specifically, SwitchTT exploits switch transformer model architecture for multi-task policy learning, allowing us to improve model capacity without proportional computation cost. Also, SwitchTT approximates the distribution rather than the expectation of trajectory value, mitigating the effects of the Monte-Carlo Value estimator suffering from poor sample complexity, especially in the sparse-reward setting. We evaluate our method using the suite of ten sparse-reward tasks from the gym-mini-grid environment.We show an improvement of 10% over Trajectory Transformer across 10-task learning and obtain up to 90% increase in offline model training speed. Our results also demonstrate the advantage of the switch transformer model for absorbing expert knowledge and the importance of value distribution in evaluating the trajectory.

Key facts. Cited by 11 papers in the map (data through 7 Oct 2026). Builds on 37 of its 50 references that are in the map. First posted on arXiv 14 Mar 2022.

Builds on

37 of 50 references.

References · 37

2017Attention Is All You Needcited by 59,932
2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
2020Language Models are Few-Shot Learnerscited by 32,199
2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
2017Proximal Policy Optimization Algorithmscited by 14,441
2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,266
2019BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehensioncited by 5,874
2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,014
2019XLNet: Generalized Autoregressive Pretraining for Language Understandingcited by 2,801
2019ALBERT: A Lite BERT for Self-supervised Learning of Language Representationscited by 2,772
2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601

25 more

2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,598
2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
2020Gradient Surgery for Multi-Task Learningcited by 1,210
2018IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architecturescited by 866
2017A Distributional Perspective on Reinforcement Learningcited by 743
2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 636
2020MOPO: Model-based Offline Policy Optimizationcited by 630
2020MOReL : Model-Based Offline Reinforcement Learningcited by 582
2013Learning Factored Representations in a Deep Mixture of Expertscited by 275
2020Decoupling Representation Learning from Reinforcement Learningcited by 263
2019Stabilizing Transformers for Reinforcement Learningcited by 262
2017Distral: Robust Multitask Reinforcement Learningcited by 231
2021MT-Opt: Continuous Multi-Task Robotic Reinforcement Learning at Scalecited by 226
2021Multi-Task Reinforcement Learning with Context-based Representationscited by 169
2022TransDreamer: Reinforcement Learning with Transformer World Modelscited by 143
2018Multi-task Deep Reinforcement Learning with PopArtcited by 138
2020Offline Reinforcement Learning from Images with Latent Space Modelscited by 94
2021Representation Matters: Offline Pretraining for Sequential Decision Makingcited by 90
2021Online and Offline Reinforcement Learning by Planning with a Learned Modelcited by 90
2021Conservative Data Sharing for Multi-Task Offline Reinforcement Learningcited by 64
2020Rewriting History with Inverse RL: Hindsight Inference for Policy Improvementcited by 59
2019Ray Interference: a Source of Plateaus in Deep Reinforcement Learningcited by 40
2021Efficient Transformers in Reinforcement Learning using Actor-Learner Distillationcited by 35
2020Knowledge Transfer in Multi-Task Deep Reinforcement Learning for Continuous Controlcited by 27

Built on by

Cited by year

Citing papers per year: 2022: 1; 2023: 5; 2024: 5.

11 papers.

2023A Survey on Transformers in Reinforcement Learningcited by 50
2023MTD-GPT: A Multi-Task Decision-Making GPT Model for Autonomous Driving at Unsignalized Intersectionscited by 38
2022On Transforming Reinforcement Learning by Transformer: The Development Trajectorycited by 35
2023Emergence of In-Context Reinforcement Learning from Noise Distillationcited by 26
2023In-Context Reinforcement Learning for Variable Action Spacescited by 23
2024Deep Generative Models for Offline Policy Learning: Tutorial, Survey, and Perspectives on Future Directionscited by 20
2024Offline Multitask Representation Learning for Reinforcement Learningcited by 14
2024HarmoDT: Harmony Multi-Task Decision Transformer for Offline Reinforcement Learningcited by 14
2023Design from Policies: Conservative Test-Time Adaptation for Offline Policy Optimizationcited by 8
2024Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learningcited by 3
2024Task-Aware Harmony Multi-Task Decision Transformer for Offline Reinforcement Learningcited by 1

Similar, not cited

Close in wording, no citation either way