Multi-Task Reinforcement Learning With Attention-Based Mixture of Experts

Multi-task learning is an important problem in reinforcement learning. Training multiple tasks together brings benefits from the shared useful information across different tasks and often achieves higher performance compared to single-task learning. However, it remains unclear how parameters in the network should be reused across tasks. Instead of naively sharing parameters across all tasks, we propose an attention-based mixture of experts multi-task reinforcement learning approach to learn a compositional policy for each task. The expert networks learn task-specific skills which specialize in different parts of multi-task representation space. To assemble the expert outputs, we propose an attention module to generate connections between tasks and experts to achieve the best performance automatically. The proposed approach can effectively learn task relationships from soft attention weights and alleviate mutual interference through independent expert networks. We use the proposed approach to improve both sample efficiency and overall performance over baseline algorithms in Meta-world, a benchmark for multi-task reinforcement learning containing 50 robotic manipulation tasks, and the multi-task MUJOCO environment, which contains four considerably different tasks with diverse state transition dynamics.