[Paper Review] Switch Trajectory Transformer with Distributional Value Approximation for Multi-Task Reinforcement Learning

Qinjie Lin, Han Liu+1 more|arXiv (Cornell University)|Mar 14, 2022

Reinforcement Learning in Robotics4 citations

TL;DR

SwitchTT proposes a multi-task offline reinforcement learning method using a sparse, expert-based Transformer architecture (Switch Transformer) and distributional value estimation to improve sample efficiency and reduce computation. It achieves 10% higher performance than Trajectory Transformer on 10 sparse-reward tasks and up to 90% faster training, demonstrating superior scalability and value estimation in multi-task settings.

ABSTRACT

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.

Motivation & Objective

• Address the challenge of high computational cost and poor sample efficiency in multi-task offline reinforcement learning with diverse, heterogeneous datasets.
• Overcome the limitations of shared policy parameters across tasks that degrade performance in multi-task settings.
• Improve value estimation in sparse-reward environments where Monte Carlo estimators suffer from poor sample complexity.
• Enable high-capacity knowledge absorption from large multi-task datasets without proportional increases in inference cost.

Proposed method

• Adopt a Switch Transformer architecture with Mixture-of-Experts (MoE) layers, where a gating network routes inputs to task-specific expert networks, enabling sparse parameter activation and reducing computation.
• Integrate three transformer heads: action prediction, dynamics modeling, and return-to-go (RTG) prediction, forming a unified sequence modeling framework for trajectory generation.
• Replace standard expected value estimation with a distributional value estimator that models the categorical distribution of trajectory returns, improving uncertainty-aware evaluation.
• Use beam search over generated trajectories, guided by the distributional value estimate, to select high-reward action sequences during planning.
• Train the model end-to-end on offline multi-task datasets, leveraging the full sequence of states, actions, returns, and task identifiers as input tokens.
• Leverage the high capacity of the Switch Transformer to absorb diverse expert demonstrations across multiple tasks while maintaining computational efficiency.

Experimental results

Research questions

• RQ1 Can a sparsely activated Transformer architecture reduce computational cost in multi-task offline RL without sacrificing model capacity or performance?
• RQ2 Does modeling the distribution of trajectory returns, rather than just the expected return, lead to more robust and accurate value estimation in sparse-reward settings?
• RQ3 How does SwitchTT compare to standard Trajectory Transformer in terms of performance and training efficiency on multi-task, sparse-reward environments?
• RQ4 To what extent can a Switch Transformer architecture effectively learn and share representations across diverse, heterogeneous tasks in offline RL?

Key findings

• SwitchTT achieves a 10% performance improvement over Trajectory Transformer across 10 sparse-reward tasks in the gym-mini-grid environment.
• The method reduces offline model training time by up to 90% compared to standard Trajectory Transformer, due to sparsity in expert computation.
• The distributional value estimator significantly improves value estimation accuracy in low-sample, sparse-reward scenarios, mitigating the poor sample complexity of Monte Carlo estimators.
• The Switch Transformer architecture enables better parameter sharing and knowledge transfer across tasks, especially when task distributions differ.
• Beam search guided by distributional value estimates leads to more reliable trajectory planning and higher cumulative rewards.
• Empirical results confirm that modeling value as a distribution, rather than an expectation, enhances policy performance and stability in multi-task offline RL.