<EMAIL_ADDRESS>
# Mixtures of Experts Unlock Parameter Scaling for Deep RL

Johan Obando-Ceron  Equal contributions  Google DeepMind  Mila - Québec AI
Institute  Université de Montréal  Ghada Sokar  Equal contributions  Google
DeepMind  Timon Willi  Equal contributions  University of Oxford  Clare Lyle
Google DeepMind  Jesse Farebrother  Google DeepMind  Mila - Québec AI
Institute  McGill University  Jakob Foerster  University of Oxford  Gintare
Karolina Dziugaite  Google DeepMind  Mila - Québec AI Institute  McGill
University  Doina Precup  Google DeepMind  Mila - Québec AI Institute  McGill
University  Pablo Samuel Castro  Google DeepMind  Mila - Québec AI Institute
Université de Montréal

###### Abstract

The recent rapid progress in (self) supervised learning models is in large
part predicted by empirical scaling laws: a model’s performance scales
proportionally to its size. Analogous scaling laws remain elusive for
reinforcement learning domains, however, where increasing the parameter count
of a model often hurts its final performance. In this paper, we demonstrate
that incorporating Mixture-of-Expert (MoE) modules, and in particular Soft
MoEs (Puigcerver et al., 2023), into value-based networks results in more
parameter-scalable models, evidenced by substantial performance increases
across a variety of training regimes and model sizes. This work thus provides
strong empirical evidence towards developing scaling laws for reinforcement
learning.

##  1 Introduction

Deep Reinforcement Learning (RL) – the combination of reinforcement learning
algorithms with deep neural networks – has proven effective at producing
agents that perform complex tasks at super-human levels (Mnih et al., 2015;
Berner et al., 2019; Vinyals et al., 2019; Fawzi et al., 2022; Bellemare et
al., 2020). While deep networks are critical to any successful application of
RL in complex environments, their design and learning dynamics in RL remain a
mystery. Indeed, recent work highlights some of the surprising phenomena that
arise when using deep networks in RL, often going against the behaviours
observed in supervised learning settings (Ostrovski et al., 2021; Kumar et
al., 2021a; Lyle et al., 2022a; Graesser et al., 2022; Nikishin et al., 2022;
Sokar et al., 2023; Ceron et al., 2023).

The supervised learning community convincingly showed that larger networks
result in improved performance, in particular for language models (Kaplan et
al., 2020). In contrast, recent work demonstrates that scaling networks in RL
is challenging and requires the use of sophisticated techniques to stabilize
learning, such as supervised auxiliary losses, distillation, and pre-training
(Farebrother et al., 2022; Taiga et al., 2022; Schwarzer et al., 2023).
Furthermore, deep RL networks are under-utilizing their parameters, which may
account for the observed difficulties in obtaining improved performance from
scale (Kumar et al., 2021a; Lyle et al., 2022a; Sokar et al., 2023). Parameter
count cannot be scaled efficiently if those parameters are not used
effectively.

Figure 1: The use of Mixture of Experts allows the performance of DQN (top)
and Rainbow (bottom) to scale with an increased number of parameters. While
Soft MoE helps in both cases and improves with scale, Top1-MoE only helps in
Rainbow, and does not improve with scale. The corresponding layer in the
baseline is scaled by the number of experts to (approximately) match
parameters. IQM scores computed over 200M environment steps over 20 games,
with 5 independent runs each, and error bars showing 95% stratified bootstrap
confidence intervals. The replay ratio is fixed to the standard $0.25$.

Architectural advances, such as transformers (Vaswani et al., 2017), adapters
(Houlsby et al., 2019), and Mixtures of Experts (MoEs; Shazeer et al., 2017),
have been central to the scaling properties of supervised learning models,
especially in natural language and computer vision problem settings. MoEs, in
particular, are crucial to scaling networks to billions (and recently
trillions) of parameters, because their modularity combines naturally with
distributed computation approaches (Fedus et al., 2022). Additionally, MoEs
induce structured sparsity in a network, and certain types of sparsity have
been shown to improve network performance (Evci et al., 2020; Gale et al.,
2019).

In this paper, we explore the effect of mixture of experts on the parameter
scalability of value-based deep RL networks, i.e., does performance increase
as we increase the number of parameters? We demonstrate that incorporating
Soft MoEs (Puigcerver et al., 2023) strongly improves the performance of
various deep RL agents, and performance improvements scale with the number of
experts used. We complement our positive results with a series of analyses
that help us understand the underlying causes for the results in Section 4.
For example, we investigate different gating mechanisms, motivating the use of
Soft MoE, as well as different tokenizations of inputs. Moreover, we analyse
different properties of the experts hidden representations, such as dormant
neurons (Sokar et al., 2023), which provide empirical evidence as to why Soft
MoE improves performance over the baseline.

Finally, we present a series of promising results that pave the way for
further research incorporating MoEs in deep RL networks in Section 5. For
instance, in Section 5.1 we show preliminary results that Soft MoE outperforms
the baseline on a set of Offline RL tasks; in Section 5.2, we evaluate Soft
MoE’s performance in low-data training regimes; and lastly, we show that
exploring different architectural designs is a fruitful direction for future
research in Section 5.3.

##  2 Preliminaries

###  2.1 Reinforcement Learning

Reinforcement learning methods are typically employed in sequential decision-
making problems, with the goal of finding optimal behavior within a given
environment (Sutton and Barto, 1998). These environments are typically
formalized as Markov Decision Processes (MDPs), defined by the tuple
$\langle\mathcal{X},\mathcal{A},\mathcal{P},\mathcal{R},\gamma\rangle$, where
$\mathcal{X}$ represents the set of states, $\mathcal{A}$ the set of available
actions,
$\mathcal{P}:\mathcal{X}\times\mathcal{A}\to\Delta(\mathcal{X})$111$\Delta(X)$
denotes a distribution over the set $X$. the transition function,
$\mathcal{R}:\mathcal{X}\times\mathcal{A}\to\mathbb{R}$ the reward function,
and $\gamma\in[0,1)$ the discount factor.

An agent’s behaviour, or policy, is expressed via a mapping
$\pi:\mathcal{X}\to\Delta(\mathcal{A})$. Given a policy $\pi$, the value
$V^{\pi}$ of a state $x$ is given by the expected sum of discounted rewards
when starting from that state and following $\pi$ from then on:
$V^{\pi}(x):=\underset{\pi,\mathcal{P}}{\mathbb{E}}\left[\sum_{t=i}^{\infty}\gamma^{t}\mathcal{R}\left(x_{t},a_{t}\right)\mid
x_{0}=x\right].$

The state-action function $Q^{\pi}$ quantifies the value of first taking
action $a$ from state $x$, and then following $\pi$ thereafter:
$Q^{\pi}(x,a):=\mathcal{R}(x,a)+\gamma\underset{x^{\prime}\sim\mathcal{P}(x,a)}{\mathbb{E}}V^{\pi}(x^{\prime})$.
Every MDP admits an optimal policy $\pi^{*}$ in the sense that
$V^{\pi^{*}}:=V^{*}\geq V^{\pi}$ uniformly over $\mathcal{X}$ for all policies
$\pi$.

Figure 2: Incorporating MoE modules into deep RL networks. Top left: Baseline
architecture; bottom left: Baseline with penultimate layer scaled up; right:
Penultimate layer replaced with an MoE module.

When $\mathcal{X}$ is very large (or infinite), function approximators are
used to express $Q$, e.g., DQN by Mnih et al. (2015) uses neural networks with
parameters $\theta$, denoted as $Q_{\theta}\approx Q$. The original
architecture used in DQN, hereafter referred to as the CNN architecture,
comprises of 3 convolutional layers followed by 2 dense layers, with ReLu
nonlinearities (Fukushima, 1969) between each pair of layers. In this work, we
mostly use the newer Impala architecture (Espeholt et al., 2018), but will
provide a comparison with the original CNN architecture. DQN also used a
replay buffer: a (finite) memory where an agent stores transitions received
while interacting with the environment, and samples mini-batches from to
compute gradient updates. The replay ratio is defined as the ratio of gradient
updates to environment interactions, and plays a role in our analyses below.

Figure 3: Tokenization types considered: PerConv (per convolution), PerFeat
(per feature), and PerSamp (per sample).

Rainbow (Hessel et al., 2018) extended the original DQN algorithm with
multiple algorithmic components to improve learning stability, sample
efficiency, and overall performance. Rainbow was shown to significantly
outperform DQN and is an important baseline in deep RL research.

###  2.2 Mixtures of Experts

Mixtures of experts (MoEs) have become central to the architectures of most
modern Large Language Models (LLMs). They consist of a set of $n$ “expert”
sub-networks activated by a gating network (typically learned and referred to
as the router), which routes each incoming token to $k$ experts (Shazeer et
al., 2017). In most cases, $k$ is smaller than the total number of experts
(k=1 in our work), thereby inducing sparser activations (Fedus et al., 2022).
These sparse activations enable faster inference and distributed computation,
which has been the main appeal for LLM training. MoE modules typically replace
the dense feed forward blocks in transformers (Vaswani et al., 2017). Their
strong empirical results have rendered MoEs a very active area of research in
the past few years (Shazeer et al., 2017; Lewis et al., 2021; Fedus et al.,
2022; Zhou et al., 2022; Puigcerver et al., 2023; Lepikhin et al., 2020; Zoph
et al., 2022; Gale et al., 2023).

Hard assignments of tokens to experts introduce a number of challenges such as
training instabilities, dropping of tokens, and difficulties in scaling the
number of experts (Fedus et al., 2022; Puigcerver et al., 2023). To address
some of these challenges, Puigcerver et al. (2023) introduced Soft MoE, which
is a fully differentiable soft assignment of tokens-to-experts, replacing
router-based hard token assignments. Soft assignment is achieved by computing
(learned) mixes of per-token weightings for each expert, and averaging their
outputs. Following the notation of Puigcerver et al. (2023), let us define the
input tokens as $\mathbf{X}\in\mathbb{R}^{m\times d}$, where $m$ is the number
of $d$-dimensional tokens. A Soft MoE layer applies a set of $n$ experts on
individual tokens,
$\left\\{f_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}\right\\}_{1:n}$. Each
expert has $p$ input- and output-slots, represented respectively by a
$d$-dimensional vector of parameters. We denote these parameters by
$\boldsymbol{\Phi}\in\mathbb{R}^{d\times(n\cdot p)}$.

Figure 4: Soft MoE yields performance gains even at high replay ratio values.
DQN (left) and Rainbow (right) with 8 experts. See Figure 6 for training
details.

The input-slots $\tilde{\mathbf{X}}\in\mathbb{R}^{(n\cdot p)\times d}$
correspond to a weighted average of all tokens:
$\tilde{\mathbf{X}}=\mathbf{D}^{\top}\mathbf{X}$, where

$\begin{gathered}\mathbf{D}_{ij}=\frac{\exp\left((\mathbf{X}\boldsymbol{\Phi})_{ij}\right)}{\sum_{i^{\prime}=1}^{m}\exp\left((\mathbf{X}\boldsymbol{\Phi})_{i^{\prime}j}\right)}.\end{gathered}$

D is typically referred to as the dispatch weights. We then denote the expert
outputs as $\tilde{\mathbf{Y}}_{i}=f_{\lfloor
i/p\rfloor}\left(\tilde{\mathbf{X}}_{i}\right)$. The output of the Soft MoE
layer $\mathbf{Y}$ is the combination of $\tilde{\mathbf{Y}}$ with the combine
weights $\mathbf{C}$ according to $\mathbf{Y}=\mathbf{C}\tilde{\mathbf{Y}}$,
where

$\begin{gathered}\mathbf{C}_{ij}=\frac{\exp\left((\mathbf{X}\boldsymbol{\Phi})_{ij}\right)}{\sum_{j^{\prime}=1}^{n\cdot
p}\exp\left((\mathbf{X}\boldsymbol{\Phi})_{ij^{\prime}}\right)}.\end{gathered}$

Note how $\mathbf{D}$ and $\mathbf{C}$ are learned only through
$\boldsymbol{\Phi}$, which we will use in our analysis. Puigcerver et al.
(2023) results suggest that Soft MoE achieves a better trade-off between
accuracy and computational cost compared to other MoE methods.

##  3  Mixture of Experts for Deep RL

Below we discuss some important design choices in our incorporation of MoE
modules into DQN-based architectures.

Figure 5: Scaling down the dimensionality of Soft MoE experts has no
significant impact on performance in Rainbow. See Figure 6 for training
details.

#### Where to place the MoEs?

Following the predominant usage of MoEs to replace dense feed-forward layers,
we replace the penultimate layer in our networks with an MoE module, where
each expert has the same dimensionality as the original dense layer. Thus, we
are effectively widening the penultimate layer’s dimensionality by a factor
equal to the number of experts. Figure 2 illustrates our deep RL MoE
architecture. We discuss some alternatives in Section 5.3.

#### What is a token?

MoEs sparsely route inputs to a set of experts (Shazeer et al., 2017) and are
mostly used in the context of transformer architectures where the inputs are
tokens (Fedus et al., 2022). For the vast majority of supervised learning
tasks where MoEs are used there is a well-defined notion of a token; except
for a few works using Transformers (Chen et al., 2021), this is not the case
for deep RL networks. Denoting by $C^{(h,w,d)}\in\mathbb{R}^{3}$ the output of
the convolutional encoders, we define _tokens_ as $d$-dimensional slices of
this output; thus, we split $C$ into $h\times w$ tokens of dimensionality $d$
(PerConv in Figure 3). This approach to tokenization is taken from what is
often used in vision tasks (see the Hybrid architecture of Dosovitskiy et al.
(2021)). We did explore other tokenization approaches (discussed in Section
4.2), but found this one to be the best performing and most intuitive.
Finally, a trainable linear projection is applied after each expert to
maintain the token size $d$ at the output.

#### What flavour of MoE to use?

We explore the top-$k$ gating architecture of Shazeer et al. (2017) following
the simplified $k=1$ strategy of Fedus et al. (2022), as well as the Soft MoE
variant proposed by Puigcerver et al. (2023); for the rest of the paper we
refer to the former as Top1-MoE and the latter as Soft MoE. We focus on these
two, as Top1-MoE is the predominantly used approach, while Soft MoE is simpler
and shows evidence of improving performance. Since Top1-MoEs activate one
expert per token while Soft MoEs can activate multiple experts per token, Soft
MoEs are arguably more directly comparable in terms of the number of
parameters when widening the dense layers of the baseline.

##  4 Empirical evaluation

Figure 6: Comparison of different tokenization methods on Rainbow with 8
experts. Soft MoE with PerConv achieves the best results. Top1-MoE works best
with PerFeat. PerSamp is worst over both architectures. See Figure 6 for
training details.

We focus our investigation on DQN (Mnih et al., 2015) and Rainbow (Hessel et
al., 2018), two value-based agents that have formed the basis of a large swath
of modern deep RL research. Recent work has demonstrated that using the ResNet
architecture (Espeholt et al., 2018) instead of the original CNN architecture
(Mnih et al., 2015) yields strong empirical improvements (Graesser et al.,
2022; Schwarzer et al., 2023), so we conduct most of our experiments with this
architecture. As in the original papers, we evaluate on 20 games from the
Arcade Learning Environment (ALE), a collection of a diverse and challenging
pixel-based environments (Bellemare et al., 2013).

Our implementation, included with this submission, is built on the Dopamine
library (Castro et al., 2018), which adds stochasticity (via sticky actions)
to the ALE (Machado et al., 2018). We use the recommendations from Agarwal et
al. (2021) for statistically-robust performance evaluations, in particular
focusing on interquartile mean (IQM). Every experiment was run for 200M
environment steps, unless reported otherwise, with 5 independent seeds, and we
report 95% stratified bootstrap confidence intervals. All experiments were run
on NVIDIA Tesla P100 GPUs, and each took on average 4 days to complete.

Figure 7: In addition to the Impala encoder, MoEs also show performance gains
for the standard CNN encoder on DQN (left) and Rainbow (right), with 8
experts. See Figure 6 for training details.

###  4.1 Soft MoE Helps Parameter Scalability

To investigate the efficacy of Top1-MoE and Soft MoE on DQN and Rainbow, we
replace the penultimate layer with the respective MoE module (see Figure 2)
and vary the number of experts. Given that each expert is a copy of the
original penultimate layer, we are effectively increasing the number of
parameters of this layer by a factor equal to the number of experts. To
compare more directly in terms of number of parameters, we evaluate simply
widening the penultimate layer of the base architectures by a factor equal to
the number of experts.

As Figure 1 demonstrates, Soft MoE provides clear performance gains, and these
gains increase with the number of experts; for instance in Rainbow, increasing
the number of experts from 1 to 8 results in a 20% performance improvement. In
contrast, the performance of the base architectures declines as we widen its
layer; for instance in Rainbow, as we increase the layer multiplier from 1 to
8 there is a performance decrease of around 40%. This is a finding consistent
with prior work demonstrating the difficulty in scaling up deep RL networks
(Farebrother et al., 2022; Taiga et al., 2022; Schwarzer et al., 2023).
Top1-MoE seems to provide gains when incorporated into Rainbow, but fails to
exhibit the parameter-scalability we observe with Soft MoE.

It is known that deep RL agents are unable to maintain performance when
scaling the replay ratio without explicit interventions (D’Oro et al., 2023;
Schwarzer et al., 2023). In Figure 4 we observe that the use of Soft MoEs
maintains a strong advantage over the baseline even at high replay ratios,
further confirming that they make RL networks more parameter efficient.

Figure 8: (Left) Evaluating the performance on 60 Atari 2600 games; (Right)
Evaluating the performance of Soft MoE with a random $\boldsymbol{\Phi}$
matrix. (Left) Even over 60 games, Soft MoE performs better than the baseline.
(Right) Learning $\boldsymbol{\Phi}$ is beneficial over a random phi,
indicating that Soft MoE’s perfomance gains are not only due to the
distribution of tokens to experts. Both plots run with Rainbow using the
Impala architecture and 8 experts. See Figure 6 for training details.

###  4.2 Impact of Design Choices

#### Number of experts

Fedus et al. (2022) argued that the number of experts is the most efficient
way to scale models in supervised learning settings. Our results in Figure 1
demonstrate that while Soft MoE does benefit from more experts, Top1-MoE does
not.

#### Dimensionality of experts

As Figure 1 confirms, scaling the corresponding layer in the base
architectures does not match the performance obtained when using Soft MoEs
(and in fact worsens it). We explored dividing the dimensionality of each
expert by the number of experts, effectively bringing the number of parameters
on-par with the original base architecture. Figure 5 demonstrates that Soft
MoE maintains its performance, even with much smaller experts. This suggests
the observed benefits come largely from the structured sparsity induced by
Soft MoEs, and not necessarily the size of each expert.

#### Gating and combining

Top1-MoEs and Soft MoEs use learned gating mechanisms, albeit different ones:
the former uses a top-1 router, selecting an expert for each token, while the
latter uses dispatch weights to assign weighted tokens to expert slots. The
learned “combiner” component takes the output of the MoE modules and combines
them to produce a single output (see Figure 2). If we use a single expert, the
only difference between the base architectures and the MoE ones are then the
extra learned parameters from the gating and combination components. Figure 1
suggests that most of the benefit of MoEs comes from the combination of the
gating/combining components with multiple experts. Interestingly, Soft MoE
with a single expert still provides performance gains for Rainbow, suggesting
that the learned $\boldsymbol{\Phi}$ matrix (see Section 2.2) has a beneficial
role to play. Indeed, the right panel of Figure 8, where we replace the
learned $\Phi$ with a random one, confirms this.

Figure 9: Load-balancing losses from Ruiz et al. (2021) are unable to improve
the performance of Top1-MoE with 8 experts on neither DQN (left) nor Rainbow
(right). See Figure 6 for training details. Figure 10: Additional analyses:
both standard Top1-MoE and Soft MoE architectures exhibit similar properties
of the hidden representation: both are more robust to dormant neurons and have
higher effective rank of both the features and the gradients. We also see that
the MoE architectures exhibit less feature norm growth than the baseline. See
Figure 6 for training details.

#### Tokenization

As mentioned above, we focus most of our investigations on PerConv
tokenization, but also explore two others: PerFeat is essentially a transpose
of PerConv, producing $d$ tokens of dimensionality $h\times w$; PerSamp uses
the entire output of the encoder as a single token (see Figure 3). In Figure 6
we can observe that while PerConv works best with Soft MoE, Top1-MoE seems to
benefit more from PerFeat.

#### Encoder

Although we have mostly used the encoder from the Impala architecture
(Espeholt et al., 2018), in Figure 7 we confirm that Soft MoE still provides
benefits when used with the standard CNN architecture from Mnih et al. (2015).

#### Game selection

To confirm that our findings are not limited to our choice of 20 games, we ran
a study over all the 60 Atari 2600 games with 5 independent seeds, similar to
previous works (Fedus et al., 2020; Ceron et al., 2023). In the left panel of
Figure 8, we observe that Soft MoE results in improved performance over the
full 60 games in the suite.

#### Number of active experts

A crucial difference between the two flavors of MoEs considered here is in the
activation of experts: Top1-MoE activates only one expert per forward pass
(hard-gating), whereas Soft MoE activates all of them. Hard-gating is known to
be a source of training difficulties, and many works have explored adding
load-balancing losses (Shazeer et al., 2017; Ruiz et al., 2021; Fedus et al.,
2022; Mustafa et al., 2022). To investigate whether Top1-MoE is
underperforming due to improper load balancing, we added the load and
importance losses proposed by Ruiz et al. (2021) (equation (7) in Appendix 2).
Figure 9 suggests the addition of these load-balancing losses is insufficient
to boost the performance of Top1-MoE. While it is possible other losses may
result in better performance, these findings suggest that RL agents benefit
from having a weighted combination of the tokens, as opposed to hard routing.

###  4.3 Additional Analysis

In the previous sections, we have shown that deep RL agents using MoE networks
are better able to take advantage of network scaling, but it is not obvious a
priori how they produce this effect. While a fine-grained analysis of the
effects of MoE modules on network optimization dynamics lies outside the scope
of this work, we zoom in on three properties known to correlate with training
instability in deep RL agents: the rank of the features (Kumar et al., 2021a),
interference between per-sample gradients (Lyle et al., 2022b), and dormant
neurons (Sokar et al., 2023). We conduct a deeper investigation into the
effect of MoE layers on learning dynamics by studying the Rainbow agent using
the Impala architecture, and using 8 experts for the runs with the MoE
modules. We track the norm of the features, the rank of the empirical neural
tangent kernel (NTK) matrix (i.e. the matrix of dot products between per-
transition gradients sampled from the replay buffer), and the number of
dormant neurons – all using a batch size of 32 – and visualize these results
in Figure 10.

We observe significant differences between the baseline architecture and the
architectures that include MoE modules. The MoE architectures both exhibit
higher numerical ranks of the empirical NTK matrices than the baseline
network, and have negligible dormant neurons and feature norms. These findings
suggest that the MoE modules have a stabilizing effect on optimization
dynamics, though we refrain from claiming a direct causal link between
improvements in these metrics and agent performance. For example, while the
rank of the ENTK is higher for the MoE agents, the best-performing agent does
not have the highest ENTK rank. The absence of pathological values of these
statistics in the MoE networks does however, suggest that whatever the precise
causal chain is, the MoE modules have a stabilizing effect on optimization.

Figure 11: Normalized performance across 17 Atari games for CQL (left) and CQL
+ C51 (right), with the ResNet (Espeholt et al., 2018) architecture and 8
experts trained on offline data. Soft MoE not only remains generally stable
with more training, but also attains higher final performance. See Figure 6
for training details.

##  5 Future directions

To provide an in-depth investigation into both the resulting performance and
probable causes for the gains observed, we focused on evaluating DQN and
Rainbow with Soft MoE and Top1-MoE on the standard ALE benchmark. The observed
performance gains suggest that these ideas would also be beneficial in other
training regimes. In this section we provide strong empirical results in a
number of different training regimes, which we hope will serve as indicators
for promising future directions of research.

Figure 12: Normalized performance across 26 Atari games for DrQ($\epsilon$)
(left) and DER (right), with the ResNet architecture (Espeholt et al., 2018)
and 8 experts. Soft MoE not only remains generally stable with more training,
but also attains higher final performance. We report interquantile mean
performance with error bars indicating 95% confidence intervals.

###  5.1 Offline RL

We begin by incorporating MoEs in offline RL, where agents are trained on a
fixed dataset without environment interactions. Following prior work (Kumar et
al., 2021b), we train the agents over 17 games and 5 seeds for 200 iterations,
where 1 iteration corresponds to 62,500 gradient updates. We evaluated on
datasets composed of 5%, 10%, 50% of the samples (drawn randomly) from the set
of all environment interactions collected by a DQN agent trained for 200M
steps (Agarwal et al., 2020). In Figure 11 we observe that combining Soft MoE
with two modern offline RL algorithms (CQL (Kumar et al., 2020) and CQL+C51
(Kumar et al., 2022)) attains the best aggregate final performance. Top1-MoE
provides performance improvements when used in conjunction with CQL+C51, but
not with CQL. Similar improvements can be observed when using 10% of the
samples (see Figure 16).

###  5.2 Agent Variants For Low-Data Regimes

DQN and Rainbow were both developed for a training regime where agents can
take millions of steps in their environments. Kaiser et al. (2020) introduced
the 100k benchmark, which evaluates agents on a much smaller data regime (100k
interactions) on 26 games. We evaluate the performance on two popular agents
for this regime: DrQ($\epsilon$), an agent based on DQN (Yarats et al., 2021;
Agarwal et al., 2021), and DER (Van Hasselt et al., 2019), an agent based on
Rainbow. When trained for 100k environment interactions we saw no real
difference between the different variants, suggesting that the benefits of
MoEs, at least in our current setup, only arise when trained for a significant
amount of interactions. However, when trained for 50M steps we do see gains in
both agents, in particular with DER (Figure 12). The gains we observe in this
setting are consistent with what we have observed so far, given that
DrQ($\epsilon$) and DER are based on DQN and Rainbow, respectively.

###  5.3 Expert Variants

Figure 13: Normalized performance over 20 games for expert variants with DQN
(left) and Rainbow (right), also investigating the use of the normalization of
Puigcerver et al. (2023). We report interquantile mean performance with
errorbars indicating 95% confidence intervals.

Our proposed MoE architecture replaces the feed-forward layer after the
encoder with an MoE, where the experts consist of a single feed-forward layer
(Figure 2). This is based on what is common practice when adding MoEs to
transformer architectures, but is by no means the only way to utilize MoEs.
Here we investigate three variants of using Soft MoE for DQN and Rainbow with
the Impala architecture: Big: Each expert is a full network.The final linear
layer is included in DQN (each expert has its own layer), but excluded from
Rainbow (so the final linear layer is shared amongst experts).222This design
choice was due to the use of C51 in Rainbow, which makes it non-trivial to
maintain‘ the “token-preserving” property of MoEs. All: A separate Soft MoE is
applied at each layer of the network, excluding the last layer for Rainbow.
Regular: the setting used in the rest of this paper. For Big and All, the
routing is applied at the the input level, and we are using PerSamp
tokenization (see Figure 3).

Puigcerver et al. (2023) propose using $\ell_{2}$ normalization to each Soft
MoE layer for large tokens, but mention that it makes little difference for
smaller tokens. Since the tokens in our experiments above are small (relative
to the large models typically using MoEs) we chose not to include this in the
experiments run thus far. This may no longer be the case with Big and All,
since we are using PerSamp tokenization on the full input; for this reason, we
investigated the impact of $\ell_{2}$ normalization when running these
results.

Figure 13 summarizes our results, from which we can draw a number of
conclusions. First, these experiments confirm our intuition that not including
regularization in the setup used in the majority of this paper performs best.
Second, consistent with Puigcerver et al. (2023), normalization seems to help
with Big experts. This is particularly noticeable in DQN, where it even
surpasses the performance of the Regular setup; the difference between the two
is most likely due to the last layer being shared (as in Rainbow) or non-
shared (as in DQN). Finally, the using separate MoE modules for each layer
(All) seems to not provide many gains, especially when coupled with
normalization, where the agents are completely unable to learn.

##  6 Related Work

#### Mixture of Experts

MoEs were first proposed by Jacobs et al. (1991) and have recently helped
scaling language models up to trillions of parameters thanks to their modular
nature, facilitating distributed training, and improved parameter efficiency
at inference (Lepikhin et al., 2020; Fedus et al., 2022). MoEs are also widely
studied in computer vision Wang et al. (2020); Yang et al. (2019); Abbas and
Andreopoulos (2020); Pavlitskaya et al. (2020), where they enable scaling
vision transformers to billions of parameters while reducing the inference
computation cost by half (Riquelme et al., 2021), help with vision-based
continual learning (Lee et al., 2019), and multi-task problems (Fan et al.,
2022). MoEs also help performance in transfer- and multi-task learning
settings, e.g. by specializing experts to sub-problems (Puigcerver et al.,
2020; Chen et al., 2023; Ye and Xu, 2023) or by addressing statistical
performance issues of routers (Hazimeh et al., 2021). There have been few
works exploring MoEs in RL for single (Ren et al., 2021; Akrour et al., 2022)
and multi-task learning (Hendawy et al., 2024). However their definition and
usage of “mixtures-of-experts” is somewhat different than ours, focusing more
on orthogonal, probabilistic, and interpretable MoEs.

#### Parameter Scalability and Efficiency in Deep RL

Lack of parameter scalability in deep RL can be partly explained by a lack of
parameter efficiency. Parameter count cannot be scaled efficiently if those
parameters are not used effectively. Recent work shows that networks in RL
under-utilize their parameters. Sokar et al. (2023) demonstrates that networks
suffer from an increasing number of inactive neurons throughout online
training. Similar behavior is observed by Gulcehre et al. (2022) in offline
RL. Arnob et al. (2021) shows that 95% of the network parameters can be pruned
at initialization in offline RL without loss in performance. Numerous works
demonstrate that periodic resetting of the network weights improves
performance (Nikishin et al., 2022; Dohare et al., 2021; Sokar et al., 2023;
D’Oro et al., 2022; Igl et al., 2020; Schwarzer et al., 2023). Another line of
research demonstrates that RL networks can be trained with a high sparsity
level ($\sim$90%) without loss in performance Tan et al. (2022); Sokar et al.
(2022); Graesser et al. (2022). These observations call for techniques to
better utilize the network parameters in RL training, such as using MoEs,
which we show decreases dormant neurons drastically over multiple tasks and
architectures. To enable scaling networks in deep RL, prior works focus on
algorithmic methods (Farebrother et al., 2022; Taiga et al., 2022; Schwarzer
et al., 2023). In contrast, we focus on alternative network topologies to
enable robustness towards scaling.

##  7 Discussion and Conclusion

As RL continues to be used for increasingly complex tasks, we will likely
require larger networks. As recent research has shown (and which our results
confirm), naïvely scaling up network parameters does not result in improved
performance. Our work shows empirically that MoEs have a beneficial effect on
the performance of value-based agents across a diverse set of training
regimes.

Mixtures of Experts induce a form of structured sparsity in neural networks,
prompting the question of whether the benefits we observe are simply a
consequence of this sparsity rather than the MoE modules themselves. Our
results suggest that it is likely a combination of both: Figure 1 demonstrates
that in Rainbow, adding an MoE module with a single expert yields
statistically significant performance improvements, while Figure 5
demonstrates that one can scale down expert dimensionality without sacrificing
performance. The right panel of Figure 8 further confirms the necessity of the
extra parameters in Soft MoE modules.

Recent findings in the literature demonstrate that while RL networks have a
natural tendency towards neuron-level sparsity which can hurt performance
(Sokar et al., 2023), they can benefit greatly from explicit parameter-level
sparsity (Graesser et al., 2022). When taken in combination with our findings,
they suggest that there is still much room for exploration, and understanding,
of the role sparsity can play in training deep RL networks, especially for
parameter scalability.

Even in the narrow setting we have focused on (replacing the penultimate layer
of value-based agents with MoEs for off-policy learning in single-task
settings), there are many open questions that can further increase the
benefits of MoEs: different values of $k$ for Top1-MoEs, different
tokenization choices, using different learning rates (and perhaps optimizers)
for routers, among others. Of course, expanding beyond the ALE could provide
more comprehensive results and insights, potentially at a fraction of the
computational expense (Ceron and Castro, 2021).

The results presented in Section 5 suggest MoEs can play a more generally
advantageous role in training deep RL agents. More broadly, our findings
confirm the impact architectural design choices can have on the ultimate
performance of RL agents. We hope our findings encourage more researchers to
further investigate this – still relatively unexplored – research direction.

#### Acknowledgements

The authors would like to thank Gheorghe Comanici, Owen He, Alex Muzio, Adrien
Ali Taïga, Rishabh Agarwal, Hugo Larochelle, and the rest of the Google
DeepMind Montreal team for valuable discussions during the preparation of this
work; Gheorghe deserves a special mention for providing us valuable feedback
on an early draft of the paper.

## References

  * Abbas and Andreopoulos (2020) A. Abbas and Y. Andreopoulos.  Biased mixtures of experts: Enabling computer vision inference under data transfer limitations.  _IEEE Transactions on Image Processing_ , 29:7656–7667, 2020.
  * Agarwal et al. (2020) R. Agarwal, D. Schuurmans, and M. Norouzi.  An optimistic perspective on offline reinforcement learning.  In H. D. III and A. Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 104–114. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/agarwal20c.html.
  * Agarwal et al. (2021) R. Agarwal, M. Schwarzer, P. S. Castro, A. C. Courville, and M. Bellemare.  Deep reinforcement learning at the edge of the statistical precipice.  _Advances in neural information processing systems_ , 34:29304–29320, 2021.
  * Akrour et al. (2022) R. Akrour, D. Tateo, and J. Peters.  Continuous action reinforcement learning from a mixture of interpretable experts.  _IEEE Trans. Pattern Anal. Mach. Intell._ , 44(10):6795–6806, oct 2022.  ISSN 0162-8828.  10.1109/TPAMI.2021.3103132.  URL https://doi.org/10.1109/TPAMI.2021.3103132.
  * Arnob et al. (2021) S. Y. Arnob, R. Ohib, S. Plis, and D. Precup.  Single-shot pruning for offline reinforcement learning.  _arXiv preprint arXiv:2112.15579_ , 2021.
  * Bellemare et al. (2013) M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 47:253–279, 2013.
  * Bellemare et al. (2020) M. G. Bellemare, S. Candido, P. S. Castro, J. Gong, M. C. Machado, S. Moitra, S. S. Ponda, and Z. Wang.  Autonomous navigation of stratospheric balloons using reinforcement learning.  _Nature_ , 588:77 – 82, 2020.
  * Berner et al. (2019) C. Berner, G. Brockman, B. Chan, V. Cheung, P. Debiak, C. Dennison, D. Farhi, Q. Fischer, S. Hashme, C. Hesse, et al.  Dota 2 with large scale deep reinforcement learning.  _arXiv preprint arXiv:1912.06680_ , 2019.
  * Castro et al. (2018) P. S. Castro, S. Moitra, C. Gelada, S. Kumar, and M. G. Bellemare.  Dopamine: A Research Framework for Deep Reinforcement Learning.  2018\.  URL http://arxiv.org/abs/1812.06110.
  * Ceron and Castro (2021) J. S. O. Ceron and P. S. Castro.  Revisiting rainbow: Promoting more insightful and inclusive deep reinforcement learning research.  In _International Conference on Machine Learning_ , pages 1373–1383. PMLR, 2021.
  * Ceron et al. (2023) J. S. O. Ceron, M. G. Bellemare, and P. S. Castro.  Small batch deep reinforcement learning.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=wPqEvmwFEh.
  * Chen et al. (2021) L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  _arXiv preprint arXiv:2106.01345_ , 2021.
  * Chen et al. (2023) Z. Chen, Y. Shen, M. Ding, Z. Chen, H. Zhao, E. G. Learned-Miller, and C. Gan.  Mod-squad: Designing mixtures of experts as modular multi-task learners.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 11828–11837, 2023.
  * Dohare et al. (2021) S. Dohare, R. S. Sutton, and A. R. Mahmood.  Continual backprop: Stochastic gradient descent with persistent randomness.  _arXiv preprint arXiv:2108.06325_ , 2021.
  * D’Oro et al. (2022) P. D’Oro, M. Schwarzer, E. Nikishin, P.-L. Bacon, M. G. Bellemare, and A. Courville.  Sample-efficient reinforcement learning by breaking the replay ratio barrier.  In _Deep Reinforcement Learning Workshop NeurIPS 2022_ , 2022.  URL https://openreview.net/forum?id=4GBGwVIEYJ.
  * D’Oro et al. (2023) P. D’Oro, M. Schwarzer, E. Nikishin, P.-L. Bacon, M. G. Bellemare, and A. Courville.  Sample-efficient reinforcement learning by breaking the replay ratio barrier.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=OpC-9aBBVJe.
  * Dosovitskiy et al. (2021) A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Espeholt et al. (2018) L. Espeholt, H. Soyer, R. Munos, K. Simonyan, V. Mnih, T. Ward, Y. Doron, V. Firoiu, T. Harley, I. Dunning, et al.  Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures.  In _Proceedings of the International Conference on Machine Learning (ICML)_ , 2018.
  * Evci et al. (2020) U. Evci, T. Gale, J. Menick, P. S. Castro, and E. Elsen.  Rigging the lottery: Making all tickets winners.  In H. D. III and A. Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 2943–2952. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/evci20a.html.
  * Fan et al. (2022) Z. Fan, R. Sarkar, Z. Jiang, T. Chen, K. Zou, Y. Cheng, C. Hao, Z. Wang, et al.  M3vit: Mixture-of-experts vision transformer for efficient multi-task learning with model-accelerator co-design.  _Advances in Neural Information Processing Systems_ , 35:28441–28457, 2022.
  * Farebrother et al. (2022) J. Farebrother, J. Greaves, R. Agarwal, C. Le Lan, R. Goroshin, P. S. Castro, and M. G. Bellemare.  Proto-value networks: Scaling representation learning with auxiliary tasks.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Fawzi et al. (2022) A. Fawzi, M. Balog, A. Huang, T. Hubert, B. Romera-Paredes, M. Barekatain, A. Novikov, F. J. R Ruiz, J. Schrittwieser, G. Swirszcz, et al.  Discovering faster matrix multiplication algorithms with reinforcement learning.  _Nature_ , 610(7930):47–53, 2022.
  * Fedus et al. (2020) W. Fedus, P. Ramachandran, R. Agarwal, Y. Bengio, H. Larochelle, M. Rowland, and W. Dabney.  Revisiting fundamentals of experience replay.  In _International Conference on Machine Learning_ , pages 3061–3071. PMLR, 2020.
  * Fedus et al. (2022) W. Fedus, B. Zoph, and N. Shazeer.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  _The Journal of Machine Learning Research_ , 23(1):5232–5270, 2022.
  * Fukushima (1969) K. Fukushima.  Visual feature extraction by a multilayered network of analog threshold elements.  _IEEE Trans. Syst. Sci. Cybern._ , 5(4):322–333, 1969.  10.1109/TSSC.1969.300225.  URL https://doi.org/10.1109/TSSC.1969.300225.
  * Gale et al. (2019) T. Gale, E. Elsen, and S. Hooker.  The state of sparsity in deep neural networks.  _CoRR_ , abs/1902.09574, 2019.  URL http://arxiv.org/abs/1902.09574.
  * Gale et al. (2023) T. Gale, D. Narayanan, C. Young, and M. Zaharia.  MegaBlocks: Efficient Sparse Training with Mixture-of-Experts.  _Proceedings of Machine Learning and Systems_ , 5, 2023.
  * Graesser et al. (2022) L. Graesser, U. Evci, E. Elsen, and P. S. Castro.  The state of sparse training in deep reinforcement learning.  In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 7766–7792. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/graesser22a.html.
  * Gulcehre et al. (2022) C. Gulcehre, S. Srinivasan, J. Sygnowski, G. Ostrovski, M. Farajtabar, M. Hoffman, R. Pascanu, and A. Doucet.  An empirical study of implicit regularization in deep offline rl.  _arXiv preprint arXiv:2207.02099_ , 2022.
  * Hazimeh et al. (2021) H. Hazimeh, Z. Zhao, A. Chowdhery, M. Sathiamoorthy, Y. Chen, R. Mazumder, L. Hong, and E. Chi.  Dselect-k: Differentiable selection in the mixture of experts with applications to multi-task learning.  _Advances in Neural Information Processing Systems_ , 34:29335–29347, 2021.
  * Hendawy et al. (2024) A. Hendawy, J. Peters, and C. D’Eramo.  Multi-task reinforcement learning with mixture of orthogonal experts.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=aZH1dM3GOX.
  * Hessel et al. (2018) M. Hessel, J. Modayil, H. V. Hasselt, T. Schaul, G. Ostrovski, W. Dabney, D. Horgan, B. Piot, M. G. Azar, and D. Silver.  Rainbow: Combining improvements in deep reinforcement learning.  In _AAAI_ , 2018.
  * Houlsby et al. (2019) N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, Q. De Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly.  Parameter-efficient transfer learning for NLP.  In K. Chaudhuri and R. Salakhutdinov, editors, _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pages 2790–2799. PMLR, 09–15 Jun 2019.  URL https://proceedings.mlr.press/v97/houlsby19a.html.
  * Igl et al. (2020) M. Igl, G. Farquhar, J. Luketina, W. Boehmer, and S. Whiteson.  Transient non-stationarity and generalisation in deep reinforcement learning.  In _International Conference on Learning Representations_ , 2020.
  * Jacobs et al. (1991) R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton.  Adaptive mixtures of local experts.  _Neural computation_ , 3(1):79–87, 1991.
  * Kaiser et al. (2020) L. Kaiser, M. Babaeizadeh, P. Miłos, B. Osiński, R. H. Campbell, K. Czechowski, D. Erhan, C. Finn, P. Kozakowski, S. Levine, A. Mohiuddin, R. Sepassi, G. Tucker, and H. Michalewski.  Model based reinforcement learning for atari.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=S1xCPJHtDB.
  * Kaplan et al. (2020) J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Kumar et al. (2020) A. Kumar, A. Zhou, G. Tucker, and S. Levine.  Conservative q-learning for offline reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33:1179–1191, 2020.
  * Kumar et al. (2021a) A. Kumar, R. Agarwal, D. Ghosh, and S. Levine.  Implicit under-parameterization inhibits data-efficient deep reinforcement learning.  In _International Conference on Learning Representations_ , 2021a.  URL https://openreview.net/forum?id=O9bnihsFfXU.
  * Kumar et al. (2021b) A. Kumar, R. Agarwal, T. Ma, A. Courville, G. Tucker, and S. Levine.  Dr3: Value-based deep reinforcement learning requires explicit regularization.  In _International Conference on Learning Representations_ , 2021b.
  * Kumar et al. (2022) A. Kumar, R. Agarwal, X. Geng, G. Tucker, and S. Levine.  Offline q-learning on diverse multi-task data both scales and generalizes.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Lee et al. (2019) S. Lee, J. Ha, D. Zhang, and G. Kim.  A neural dirichlet process mixture model for task-free continual learning.  In _International Conference on Learning Representations_ , 2019.
  * Lepikhin et al. (2020) D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen.  Gshard: Scaling giant models with conditional computation and automatic sharding.  In _International Conference on Learning Representations_ , 2020.
  * Lewis et al. (2021) M. Lewis, S. Bhosale, T. Dettmers, N. Goyal, and L. Zettlemoyer.  Base layers: Simplifying training of large, sparse models.  In _International Conference on Machine Learning_ , 2021.  URL https://api.semanticscholar.org/CorpusID:232428341.
  * Lyle et al. (2022a) C. Lyle, M. Rowland, and W. Dabney.  Understanding and preventing capacity loss in reinforcement learning.  In _International Conference on Learning Representations_ , 2022a.  URL https://openreview.net/forum?id=ZkC8wKoLbQ7.
  * Lyle et al. (2022b) C. Lyle, M. Rowland, W. Dabney, M. Kwiatkowska, and Y. Gal.  Learning dynamics and generalization in deep reinforcement learning.  In _International Conference on Machine Learning_ , pages 14560–14581. PMLR, 2022b.
  * Machado et al. (2018) M. C. Machado, M. G. Bellemare, E. Talvitie, J. Veness, M. Hausknecht, and M. Bowling.  Revisiting the arcade learning environment: evaluation protocols and open problems for general agents.  _J. Artif. Int. Res._ , 61(1):523–562, jan 2018\.  ISSN 1076-9757.
  * Mnih et al. (2015) V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, S. Petersen, C. Beattie, A. Sadik, I. Antonoglou, H. King, D. Kumaran, D. Wierstra, S. Legg, and D. Hassabis.  Human-level control through deep reinforcement learning.  _Nature_ , 518(7540):529–533, Feb. 2015.
  * Mustafa et al. (2022) B. Mustafa, C. Riquelme, J. Puigcerver, R. Jenatton, and N. Houlsby.  Multimodal contrastive learning with limoe: the language-image mixture of experts.  _Advances in Neural Information Processing Systems_ , 35:9564–9576, 2022.
  * Nikishin et al. (2022) E. Nikishin, M. Schwarzer, P. D’Oro, P.-L. Bacon, and A. Courville.  The primacy bias in deep reinforcement learning.  In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 16828–16847. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/nikishin22a.html.
  * Ostrovski et al. (2021) G. Ostrovski, P. S. Castro, and W. Dabney.  The difficulty of passive learning in deep reinforcement learning.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=nPHA8fGicZk.
  * Pavlitskaya et al. (2020) S. Pavlitskaya, C. Hubschneider, M. Weber, R. Moritz, F. Huger, P. Schlicht, and M. Zollner.  Using mixture of expert models to gain insights into semantic segmentation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops_ , pages 342–343, 2020.
  * Puigcerver et al. (2020) J. Puigcerver, C. R. Ruiz, B. Mustafa, C. Renggli, A. S. Pinto, S. Gelly, D. Keysers, and N. Houlsby.  Scalable transfer learning with expert models.  In _International Conference on Learning Representations_ , 2020.
  * Puigcerver et al. (2023) J. Puigcerver, C. Riquelme, B. Mustafa, and N. Houlsby.  From sparse to soft mixtures of experts, 2023.
  * Ren et al. (2021) J. Ren, Y. Li, Z. Ding, W. Pan, and H. Dong.  Probabilistic mixture-of-experts for efficient deep reinforcement learning.  _CoRR_ , abs/2104.09122, 2021.  URL https://arxiv.org/abs/2104.09122.
  * Riquelme et al. (2021) C. Riquelme, J. Puigcerver, B. Mustafa, M. Neumann, R. Jenatton, A. Susano Pinto, D. Keysers, and N. Houlsby.  Scaling vision with sparse mixture of experts.  _Advances in Neural Information Processing Systems_ , 34:8583–8595, 2021.
  * Ruiz et al. (2021) C. R. Ruiz, J. Puigcerver, B. Mustafa, M. Neumann, R. Jenatton, A. S. Pinto, D. Keysers, and N. Houlsby.  Scaling vision with sparse mixture of experts.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=FrIDgjDOH1u.
  * Schwarzer et al. (2023) M. Schwarzer, J. S. Obando Ceron, A. Courville, M. G. Bellemare, R. Agarwal, and P. S. Castro.  Bigger, better, faster: Human-level Atari with human-level efficiency.  In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 30365–30380. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/schwarzer23a.html.
  * Shazeer et al. (2017) N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=B1ckMDqlg.
  * Sokar et al. (2022) G. Sokar, E. Mocanu, D. C. Mocanu, M. Pechenizkiy, and P. Stone.  Dynamic sparse training for deep reinforcement learning.  In _International Joint Conference on Artificial Intelligence_ , 2022\.
  * Sokar et al. (2023) G. Sokar, R. Agarwal, P. S. Castro, and U. Evci.  The dormant neuron phenomenon in deep reinforcement learning.  In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 32145–32168. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/sokar23a.html.
  * Sutton and Barto (1998) R. S. Sutton and A. G. Barto.  _Introduction to Reinforcement Learning_.  MIT Press, Cambridge, MA, USA, 1st edition, 1998.  ISBN 0262193981.
  * Taiga et al. (2022) A. A. Taiga, R. Agarwal, J. Farebrother, A. Courville, and M. G. Bellemare.  Investigating multi-task pretraining and generalization in reinforcement learning.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Tan et al. (2022) Y. Tan, P. Hu, L. Pan, J. Huang, and L. Huang.  Rlx2: Training a sparse deep reinforcement learning model from scratch.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Van Hasselt et al. (2019) H. P. Van Hasselt, M. Hessel, and J. Aslanides.  When to use parametric models in reinforcement learning?  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin.  Attention is all you need.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.
  * Vinyals et al. (2019) O. Vinyals, I. Babuschkin, W. M. Czarnecki, M. Mathieu, A. Dudzik, J. Chung, D. H. Choi, R. Powell, T. Ewalds, P. Georgiev, et al.  Grandmaster level in starcraft ii using multi-agent reinforcement learning.  _Nature_ , 575(7782):350–354, 2019.
  * Wang et al. (2020) X. Wang, F. Yu, L. Dunlap, Y.-A. Ma, R. Wang, A. Mirhoseini, T. Darrell, and J. E. Gonzalez.  Deep mixture of experts via shallow embedding.  In _Uncertainty in artificial intelligence_ , pages 552–562. PMLR, 2020.
  * Yang et al. (2019) B. Yang, G. Bender, Q. V. Le, and J. Ngiam.  Condconv: Conditionally parameterized convolutions for efficient inference.  _Advances in neural information processing systems_ , 32, 2019.
  * Yarats et al. (2021) D. Yarats, I. Kostrikov, and R. Fergus.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=GY6-6sTvGaf.
  * Ye and Xu (2023) H. Ye and D. Xu.  Taskexpert: Dynamically assembling multi-task representations with memorial mixture-of-experts.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 21828–21837, 2023.
  * Zhou et al. (2022) Y. Zhou, T. Lei, H. Liu, N. Du, Y. Huang, V. Zhao, A. M. Dai, Q. V. Le, J. Laudon, et al.  Mixture-of-experts with expert choice routing.  _Advances in Neural Information Processing Systems_ , 35:7103–7114, 2022.
  * Zoph et al. (2022) B. Zoph, I. Bello, S. Kumar, N. Du, Y. Huang, J. Dean, N. Shazeer, and W. Fedus.  St-moe: Designing stable and transferable sparse expert models, 2022.

##  Appendix A Extra results

Figure 14: Results for architectural ablations as described in Section 5.3 on
DQN. Additionally, we investigate the effect of the normalization that was
proposed in the original Soft MoE paper. Figure 15: Results for architectural
exploration as described in Section 5.3 on Rainbow. Additionally, we
investigate the effect of the normalization that was proposed in the original
Soft MoE paper.

Figure 16: Normalized performance across 17 Atari games for CQL+C51. x-axis
represents gradient steps; no new data is collected. Right: 10% and Right: 50%
uniform replay. We report IQM with 95% stratified bootstrap CIs (Agarwal et
al., 2021)

.