Mixture of Experts in a Mixture of RL settings (2406.18420v1)

Abstract: Mixtures of Experts (MoEs) have gained prominence in (self-)supervised learning due to their enhanced inference efficiency, adaptability to distributed training, and modularity. Previous research has illustrated that MoEs can significantly boost Deep Reinforcement Learning (DRL) performance by expanding the network's parameter count while reducing dormant neurons, thereby enhancing the model's learning capacity and ability to deal with non-stationarity. In this work, we shed more light on MoEs' ability to deal with non-stationarity and investigate MoEs in DRL settings with "amplified" non-stationarity via multi-task training, providing further evidence that MoEs improve learning capacity. In contrast to previous work, our multi-task results allow us to better understand the underlying causes for the beneficial effect of MoE in DRL training, the impact of the various MoE components, and insights into how best to incorporate them in actor-critic-based DRL networks. Finally, we also confirm results from previous work.

• The paper reveals that the 'Big' MoE architecture significantly outperforms others in mitigating non-stationarity in multi-task and continual reinforcement learning settings.
• It demonstrates that strategic placement of MoE modules enhances network plasticity and reduces dormant neurons in actor networks.
• Experimental evaluations on games like SpaceInvaders, Breakout, and Asterix validate the efficacy of flexible routing methods in DRL.

Mixture of Experts in a Mixture of RL Settings: A Technical Perspective

The paper "Mixture of Experts in a Mixture of RL Settings" presents a comprehensive exploration of the efficacy of Mixture of Experts (MoEs) applied to Deep Reinforcement Learning (DRL) in highly non-stationary environments. It explores two specific settings: Multi-Task Reinforcement Learning (MTRL) and Continual Reinforcement Learning (CRL). The motivation behind the paper is rooted in prior evidence that MoEs can enhance DRL performance by scaling networks efficiently while mitigating the issue of dormant neurons, thus preserving network plasticity.

Key Findings and Methodologies

Experimental Setup and Architectures

The paper evaluates multiple MoE architectures in the context of DRL, focusing on how these architectures can be adapted to handle the amplified non-stationarity inherent in MTRL and CRL settings. The distinct architectures under consideration include:

1. Middle: MoE modules replace the penultimate layer.
2. Final: The final layer is replaced by MoE modules.
3. All: All layers are substituted with MoE modules.
4. Big: The network contains a single MoE module, with each expert comprising the entire original network.

These architectures are evaluated using the PureJaxRL codebase and optimised MinAtar environments, with experiments encompassing three different games: SpaceInvaders, Breakout, and Asterix.

Strong Numerical Results

The Big MoE architecture stands out consistently across various metrics:

• In the MTRL setting, Big significantly outperforms others, suggesting that segregating tasks effectively, each handled by a separate expert network, yields superior performance.
• In the CRL setting, where the order and infrequent switching of tasks induce high non-stationarity, Big-Hardcoded shows the strongest retention and performance, which is critical for continual learning across tasks.

Interestingly, the results exposed limitations in other architectures like All, where suboptimal hyperparameters might have hindered performance.

Routing Mechanisms

The paper rigorously investigates the routing strategies for MoEs:

• SoftMoE and Hardcoded routers perform effectively in MTRL settings but show mixed results in CRL due to discontinuities in environment dynamics.
• TopK routing often struggles, highlighting the importance of flexible, yet robust, routing strategies under different DRL scenarios.

Adding task identifiers or gradient information to the routing strategies did not yield significant improvements, challenging some existing hypotheses from supervised learning contexts.

Theoretical and Practical Implications

Network Plasticity and Dormant Neurons

A significant insight from the research is the MoE's ability to decrease dormant neurons across different architectures. This supports the premise that MoEs maintain network plasticity, crucial for environments with high levels of non-stationarity, such as MTRL and CRL.

Expert Specialisation

The results demonstrate a tendency toward expert specialisation within MoE architectures without explicit load-balancing. This specialisation is critical in DRL where different tasks or environments might require vastly different policy structures.

Actor-Critic Networks

The findings also pointed out that MoEs have a more pronounced impact on actor networks compared to critic networks, which aligns with the understanding that actor networks can better handle higher sparsity levels while maintaining performance.

Speculation on Future Developments

The application of MoEs in DRL, especially under non-stationary settings, opens up several avenues for future research and development:

1. Task Curricula and Overfitting: Investigating optimal task curricula could significantly enhance policy retention and overall performance. Understanding when and how task overfitting occurs can lead to better algorithm design.
2. Multi-Agent Systems: Extending MoE use to multi-agent DRL, where experts could represent different agents in cooperative or competitive settings, can leverage the scalability and flexibility of MoEs.
3. Fine-Tuning Hyperparameters: More focused hyperparameter tuning for architectures like All could uncover performance potentials currently obscured by suboptimal settings.

In conclusion, this paper provides robust evidence supporting the use of MoEs in enhancing DRL in non-stationary environments. The insights into network plasticity, expert specialisation, and the impact of different routing strategies offer valuable directions for both theoretical advancements and practical applications in reinforcement learning.