# A Pytorch Reproduction of Masked Generative Image Transformer

Victor Besnier
Valeo.ai, Prague Mickael Chen
Valeo.ai, Paris

###### Abstract

In this technical report, we present a reproduction of MaskGIT: Masked
Generative Image Transformer [3], using PyTorch [11]. The approach involves
leveraging a masked bidirectional transformer architecture, enabling image
generation with only few steps ($8\sim 16$ steps) for $512\times 512$
resolution images, i.e., $\sim$64x faster than an auto-regressive approach.
Through rigorous experimentation and optimization, we achieved results that
closely align with the findings presented in the original paper. We match the
reported FID of 7.32 with our replication and obtain 7.59 with similar
hyperparameters on ImageNet at resolution $512\times 512$. Moreover, we
improve over the official implementation with some minor hyperparameter
tweaking, achieving FID of 7.26. At the lower resolution of $256\times 256$
pixels, our reimplementation scores 6.80, in comparison to the original
paper’s 6.18. To promote further research on Masked Generative Models and
facilitate their reproducibility, we released our code and pre-trained weights
openly at https://github.com/valeoai/MaskGIT-pytorch/.

Figure 1: Examples generated on ImageNet at $512\times 512$ demonstrate the
effectiveness of our reproduction. Hyperparameters for this set of examples
are Gumbel temperature set to 7, Softmax temperature set to 1.3, CFG set to 9,
scheduler set to arccos, and scheduler step set to 32

##  1 Masked Generative Models

In recent years, advancements in image synthesis have exhibited significant
progress. Notably, diffusion models [5] have gained considerable traction,
surpassing GANs in popularity due to its ease of training and capacity to
capture a more diversified distribution. But, this paper directs its focus
towards another type of generative methods known as Masked Generative models
(MGM) [2, 3, 10, 13, 15, 16]. The fundamental concept involves harnessing the
capabilities of a VQGAN [6] to discretize and reduce an image into visual
tokens. Once this discretization is achieved, a bidirectional transformer [4]
is trained to unmask a randomly selected set of visual tokens. During
inference, the transformer progressively unmasks a fully masked image. This
process makes it possible to unmask tokens in parallel. More information can
be found in [3].

Conceptually, MGMs have strong connections to several popular and successful
deep learning methods. They are very close to Masked Auto-Encoders [7], a
strong family of models for self-supervised learning, with whom they share
similar transformers architectures and the masked reconstruction training
objective, except in a token space. MGMs are also connected to auto-regressive
models [12, 1], a very strong paradigm for data generation, as both perform
generation by iteratively predicting tokens, but only MGMs are able to
parallelize the process. Finally, MGMs can be linked to discrete reverse
diffusion, when the forward noise process is replaced by a discrete patch
dropping scheme. The sampling process is progressively retrieving patches over
multiple iterations. As MGMs are in very close proximity to those tried and
tested methods, we argue that they are very promising research venues that
deserve more attention than they currently have. This idea is further
supported by the fact that a Masked Generative Model, MAGVIT-v2 [16], reached
state-of-the-art results for both text to image and text to video generation.

However, despite its promises, it has to be noted that very few public
releases of MGMs have been made available so far. At the time of writing, most
prominent papers for image generation do not have full release of code and
pretrained weights, if any at all. For instance, the founding paper, MaskGIT
[3], comes with pretrained model and inference
script111https://github.com/google-research/maskgit/, but no training nor
evaluation scripts. A text-to-image model, MUSE [2], has no public code
release at all. The MAGVIT [15], introduced for video generation, can also be
used for images and do provide code, but no pretrained weights. For video
prediction, TECO [14] has both code and pretrained weights, but nothing
trained on image datasets, and Phenaki [13] has no official code release at
all. Moreover, when existing, all these codebases are JAX implementations, and
very few resources for Mask Generative Models exist in Pytorch. Still, on-
going reproduction efforts of MUSE222https://github.com/lucidrains/muse-
maskgit-pytorch and Phenaki333https://github.com/lucidrains/phenaki-pytorch in
Pytorch should be noted but they have yet to publicize results or pretrained
models. To the best of our knowledge, the only complete image Masked
Generative Model released in Pytorch is that of MAGE [10], that proposed a
non-conditional generative model where MaskGIT accepts class token that allows
for conditioning.

In this context, this technical report sets out to reproduce the findings of
the original MaskGIT paper [3], as their training code is not available and
because no convincing reproduction has been publicly released so far. With
this report, we provide all the code necessary for replicating the results,
including training code and our pretrained models. This work aims at helping
researchers to explore the topic of Masked Generative Modeling.

We organize this report as follows: in Section 2, we succinctly describe the
model, training and sampling techniques, and discuss differences between our
reproduction and the details in the paper or official implementation. In
Section 3, we present both quantitative and qualitative results obtained from
our networks and, in particular, we ablate the different sampling
hyperparameters.

##  2 Implementation Details

We utilized a pre-trained VQGAN consisting of 72.142M parameters, trained
exclusively on the ImageNet dataset. This VQGAN employs a codebook of 1024
codes and reduces a $256\times 256$ (resp. $512\times 512$) image to a
$16\times 16$ (resp. $32\times 32$) token representation. For additional
information, readers can refer to the official
repository444https://github.com/CompVis/taming-transformers.

The bidirectional transformer that generates the visual tokens consists of
174.161M parameters on ImageNet $256\times 256$ and 176.307M parameters on
ImageNet $512\times 512$. This sums up to 246.30M and 248.44M respectively,
including the VQGAN. Hyperparameters settings are shown in Table 1.

Table 1: Transformer architecture Hidden Dimension | Codebook Size | Depth | Attention Heads | MLP Size | Dropout
---|---|---|---|---|---
768 | 1024 | 24 | 16 | 3072 | 0.1

For the inputs, we concatenate the conditional tokens with the visual tokens,
i.e., the model takes as input ($16\times 16$) + 1 tokens for ImageNet
$256\times 256$, and ($32\times 32$) + 1 tokens for ImageNet $512\times 512$.
The model is trained for token unmasking, using cross-entropy loss with label
smoothing of 0.1. The optimizer employed is AdamW with a learning rate of
$1e^{-4}$, betas=(0.9, 0.96), and a weight decay of $1e^{-5}$. We utilize an
arccos scheduler for masking during training, regardless of the image
resolution. Additionally, we drop 10% of the conditions for classifier-free
guidance. We train our methods on both ImageNet $256\times 256$ and $512\times
512$. This well-known classification dataset consists of 1.2M images
categorized into 1000 classes. For data augmentation, we simply use random
cropping and horizontal flipping.

During the training of the masked transformer, a batch size of 512 was
employed over 300 epochs, utilizing 8 GPUs (768 GPU hours on Nvidia A100) for
755,200 iterations. Subsequently, we fine-tuned the same model for
$\sim$750,000 additional iterations on ImageNet $512\times 512$ with a batch
size of 128 and 384 GPU hours on Nvidia A100. In total, this project,
including training, testing and debugging, consumed $\sim$3,500 GPU hours on
Nvidia A100.

While a lot of information to reproduce the results was available in the main
paper, some points remain uncertain. In particular, we discuss here the main
differences between the paper, the official github implementation, and our own
design choices:

  * •

For the architecture of the unmasking transformer, we try to stay as close as
possible to the JAX implementation. But it is still unclear why the class
embeddings and the visual embeddings are shared in a single layer. This oddity
results in implementing the token classification head as a dot product between
the outputs of the transformer and each of the embeddings (classes and visual
tokens), keeping only the similarity with visual embeddings to compute the
cross-entropy. The remaining 1001 values that correspond to the similarity
with the class embeddings are dropped in the inference code, and we don’t know
if they were actually used for training in the official release.

  * •

For sampling, the Gumbel noise injected in the confidence is not mentioned in
the paper but taken from the official implementation.

  * •

In addition, and not documented in the official releases, we used classifier-
free guidance, and chose a learning rate of $1e^{-4}$.

  * •

We also use a different number of inference steps (15 steps) to achieve higher
performance when compared to the original paper (12 steps).

##  3 Performance on ImageNet

###  3.1 Quantitative Results

Using the hyperparameters for sampling presented in Table 2, we report our
quantitative results in Table 3, as well as those from the original paper.
Overall, our reproduction shows very similar results for both ImageNet
$256\times 256$ and $512\times 512$. Slight differences are that our
reproduction emphasizes better quality (Inception Score , Precision) and
matches the diversity (Fréchet Inception Distance, Recall). While only 8 steps
are sufficient to generate good images in the lower resolution, one must
increase the number of decoding steps up to 15 to achieve the best performance
for ImageNet $512\times 512$. Moreover, a higher Gumbel noise is required to
maintain diversity at higher resolutions (see also Figure 2(b)). Finally, it
is noteworthy that the best performance is achieved using the arccos
scheduler, but results vary according to the number of steps, more information
in Table 4.

Table 2: Hyperparameter used to get the best FID score. Resolution | Softmax Temp. | Gumbel Temp. | CFG (w) | Schedule | Step
---|---|---|---|---|---
$256\times 256$ | 1.0 | 4.5 | 3.0 | arccos | 8
$512\times 512$ | 1.0 | 7.0 | 2.8 | arccos | 15
Table 3: Comparison between this work (Ours) and the official paper on ImageNet 256 and 512. Metric | Ours |  MaskGIT [3] | Ours |  MaskGIT [3]
---|---|---|---|---
| ($256\times 256$) | ($256\times 256$) | ($512\times 512$) | ($512\times 512$)
FID | 6.80 | 6.18 | 7.26 | 7.32
IS | 214.0 | 182.1 | 223.0 | 156.0
Precision | 0.82 | 0.80 | 0.85 | 0.78
Recall | 0.51 | 0.51 | 0.49 | 0.50
Density | 1.25 | - | 1.33 | -
Coverage | 0.84 | - | 0.86 | -

In Figure 2, we present an ablation study on various hyperparameters related
to the sampling process. For each curve or table representing an ablation
study, all hyperparameters except the one being tested are held constant, as
outlined below: scheduler is set to arccos, number of step is set to 12,
softmax temperature is set to 1.0, cfg weight is set to 3.0 and Gumbel
temperature: 7.

Figure 2(a) show that the Inception Score (IS) exhibits a consistent increase
throughout the training process, while the Fréchet Inception Distance (FID)
demonstrates a corresponding decrease. This curve suggests that prolonging the
training could potentially enhance the performance significantly, particularly
for ImageNet $256\times 256$.

Furthermore, we find out that injecting noise in the confidence score
substantially improves the diversity of generated samples. Indeed, as can be
seen in Figure 2(b), introducing Gumbel noise in the confidence score, with a
linear decay over the sampling, significantly improves the FID from $66.7$ to
$7.7$. This trick, taken from the official code, was not discussed in the
original paper.

Additionally, we study the influence of the number of steps on the quality of
the samples. As expected, we show in Figure 2(c) that as the number of steps
increases, the sampling quality improves. However, it is notable that the
performance saturates beyond a certain number of steps, reaching best FID at
15 steps.

Also, in Figure 2(d), we perform experiments on the classifier free guidance
(cfg) to strike a balance between fidelity, represented by FID, and quality,
as indicated by IS. An optimal trade-off seems to manifest at approximately
cfg = 3, highlighting the importance of appropriate configuration for
achieving the desired results with our model.

(a) Training FID and IS (ImageNet $256\times 256$)

(b) Gumbel noise (ImageNet $512\times 512$)

(c) Number of steps (ImageNet $512\times 512$)

(d) Classifier Free Guidance weight ($512\times 512$)

Figure 2: Hyperparameters search: ablation on crucial parameters to control
sampling quality and diversity: number of training epochs, Gumbel noise,
number of steps and the cfg weight.

Finally, in Table 4, we conduct a comparative analysis of various schedulers
at $512\times 512$ pixel resolution. Despite the model being trained with an
arccos scheduler, the results indicate that other schedulers at inference
could have some advantages. Indeed, when using 12 steps, square scheduler
significantly improves FID while maintaining a reasonably high Inception
Score. Conversely, a root scheduler yields superior quality, albeit at the
cost of a higher FID. Unfortunately, none of these schedulers scale to a
higher number of steps, and best results are obtained with the arccos
scheduler with 15 steps or more. Nevertheless, these results imply the choice
of a scheduler should change based on the number of sampling steps and compute
budget.

Table 4: Scheduler configurations and metrics on ImageNet $512\times 512$ Scheduler | FID | IS | Precision | Recall | Density | Coverage
---|---|---|---|---|---|---
root | 8.21 | 252.96 | 0.8619 | 0.4644 | 1.3699 | 0.8666
linear | 7.80 | 238.21 | 0.8547 | 0.4682 | 1.3505 | 0.8633
square | 7.50 | 224.07 | 0.8355 | 0.4904 | 1.2887 | 0.8524
cosine | 7.59 | 229.42 | 0.8423 | 0.4832 | 1.3052 | 0.8546
arccos | 7.65 | 219.32 | 0.8345 | 0.4826 | 1.2794 | 0.8437

###  3.2 Qualitative Results

Here, we analyze our visual outcomes, showcased in Figure 1, focusing on
$512\times 512$ resolution samples. These particular examples represent the
visually best among the 10 random samples for each class. Each image requires
0.9365 seconds and 32 steps to be generated. In Figure 3, we compare our
samples to those showcased in the official paper. Our results, while seemingly
displaying slightly less diversity, highlight higher quality, and are not
cherry-picked. Moreover, those images take only 0.4406 seconds and 15 steps to
be generated at $512\times 512$ resolution. And only 0.036 seconds is required
to generate a $256\times 256$ pixels image with 8 forward steps. All results
were obtained with Nvidia A100 GPUs. It is important to notice that it is
$\sim$ 64x faster than auto-regressive methods.

(a) Official Ostrich (009)

(b) Official Burger (933)

(c) Official Volcano (980)

(d) Ours Ostrich (009)

(e) Ours Burger (933)

(f) Ours Volcano (980)

Figure 3: Diversity comparison between the official paper (top row) and our
reproduction (bottom row). Without cherry pinking on our side, our methods
exhibit a little bit less diversity but a higher quality Figure 4:
Intermediate images generated at $256\times 256$ resolution. The first row
showcases the binary mask, the second row exhibits the dog generated in
association with the binary mask, and the last row presents another example of
a sailboat.

In Figure 4, we display the intermediate predictions of our models during the
sampling process. Remarkably, the model produces reasonably good images even
at an early stage of sampling.

Additionally, Figure 5 presents our inpainting results for a rooster (ImageNet
007) and a zebra (ImageNet 340) integrated into a Cityscapes image. We
demonstrate that our adaptation of MaskGIT successfully achieves filling
ImageNet classes within a road scene, showcasing the potential of our approach
in this domain.

Figure 5: Image inpainting:A rooster (ImageNet 007) and a zebra (ImageNet
340), generated and inpainted in a Cityscapes image.

##  4 Conclusion

We released a Pytorch [11] reproduction of MaskGIT [3], the founding work for
Masked Generative models. Possible next direction is to extend the repository
and add models for different modalities, like text-to-image (e.g. MUSE [2]) or
text-to-video (e.g. MAGVIT [15]). It can also incorporate further Masked
Generation improvement such as Token Critic [9] or Frequency Adaptive Sampling
[8]. It is our hope that this work can help other researchers to further
explore the capabilities of Masked Generative Modeling.

## Acknowledgments

This research received the support of EXA4MIND, a European Union’s Horizon
Europe Research and Innovation program under grant agreement N°101092944.
Views and opinions expressed are however those of the author(s) only and do
not necessarily reflect those of the European Union or the European
Commission. Neither the European Union nor the granting authority can be held
responsible for them.

## References

  * [1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  In NeurIPS, 2020.
  * [2] Huiwen Chang, Han Zhang, Jarred Barber, Aaron Maschinot, Jose Lezama, Lu Jiang, Ming-Hsuan Yang, Kevin Patrick Murphy, William T Freeman, Michael Rubinstein, et al.  Muse: Text-to-image generation via masked generative transformers.  In ICML, 2023.
  * [3] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman.  Maskgit: Masked generative image transformer.  In CVPR, June 2022.
  * [4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: pre-training of deep bidirectional transformers for language understanding.  In NAACL-HLT, 2019.
  * [5] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  In NeurIPS, 2021.
  * [6] Patrick Esser, Robin Rombach, and Björn Ommer.  Taming transformers for high-resolution image synthesis.  In CVPR, 2020.
  * [7] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick.  Masked autoencoders are scalable vision learners.  In CVPR, 2022.
  * [8] Jaewoong Lee, Sangwon Jang, Jaehyeong Jo, Jaehong Yoon, Yunji Kim, Jin-Hwa Kim, Jung-Woo Ha, and Sung Ju Hwang.  Text-conditioned sampling framework for text-to-image generation with masked generative models.  In ICCV, 2023.
  * [9] José Lezama, Huiwen Chang, Lu Jiang, and Irfan Essa.  Improved masked image generation with token-critic.  In ECCV, 2022.
  * [10] Tianhong Li, Huiwen Chang, Shlok Kumar Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan.  MAGE: masked generative encoder to unify representation learning and image synthesis.  In CVPR, 2023.
  * [11] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In NeurIPS, 2019.
  * [12] Aäron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu.  Pixel recurrent neural networks.  In Maria-Florina Balcan and Kilian Q. Weinberger, editors, ICML, JMLR Workshop and Conference Proceedings, 2016.
  * [13] Ruben Villegas, Mohammad Babaeizadeh, Pieter-Jan Kindermans, Hernan Moraldo, Han Zhang, Mohammad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan.  Phenaki: Variable length video generation from open domain textual descriptions.  In ICLR, 2023.
  * [14] Wilson Yan, Danijar Hafner, Stephen James, and Pieter Abbeel.  Temporally consistent transformers for video generation.  In ICML, 2023.
  * [15] Lijun Yu, Yong Cheng, Kihyuk Sohn, José Lezama, Han Zhang, Huiwen Chang, Alexander G. Hauptmann, Ming-Hsuan Yang, Yuan Hao, Irfan Essa, and Lu Jiang.  MAGVIT: masked generative video transformer.  In CVPR, 2023.
  * [16] Lijun Yu, José Lezama, Nitesh B. Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Agrim Gupta, Xiuye Gu, Alexander G. Hauptmann, Boqing Gong, Ming-Hsuan Yang, Irfan Essa, David A. Ross, and Lu Jiang.  Language model beats diffusion – tokenizer is key to visual generation.  CoRR, abs/2310.05737, 2023.


# Probabilistic Mixture-of-Experts for Efficient Deep Reinforcement Learning

Jie Ren  Yewen Li  Zihan Ding  Wei Pan  Hao Dong

###### Abstract

Deep reinforcement learning (DRL) has successfully solved various problems
recently, typically with a unimodal policy representation. However, grasping
distinguishable skills for some tasks with non-unique optima can be essential
for further improving its learning efficiency and performance, which may lead
to a multimodal policy represented as a mixture-of-experts (MOE). To our best
knowledge, present DRL algorithms for general utility do not deploy this
method as policy function approximators due to the potential challenge in its
differentiability for policy learning. In this work, we propose a
probabilistic mixture-of-experts (PMOE) implemented with a Gaussian mixture
model (GMM) for multimodal policy, together with a novel gradient estimator
for the indifferentiability problem, which can be applied in generic off-
policy and on-policy DRL algorithms using stochastic policies, _e.g._ , Soft
Actor-Critic (SAC) and Proximal Policy Optimisation (PPO). Experimental
results testify the advantage of our method over unimodal polices and two
different MOE methods, as well as a method of option frameworks, based on the
above two types of DRL algorithms, on six MuJoCo tasks. Different gradient
estimations for GMM like the reparameterisation trick (Gumbel-Softmax) and the
score-ratio trick are also compared with our method. We further empirically
demonstrate the distinguishable primitives learned with PMOE and show the
benefits of our method in terms of exploration.

Machine Learning, ICML

##  1 Introduction

Current deep reinforcement learning (DRL) (Sutton & Barto, 1998; Haarnoja et
al., 2018; Schulman et al., 2017; Mnih et al., 2016; Lillicrap et al., 2016;
Dong et al., 2020) methods are mostly built on parameterised models,
specifically, neural networks. Function approximation lies at the core of DRL
methods, including both value function approximation (van Hasselt et al.,
2015; Wang et al., 2015) and policy function approximation (Haarnoja et al.,
2018; Schulman et al., 2017). The methods with policy-centric function
approximation in DRL are called policy-based algorithms, which typically apply
the optimisation method: policy gradient. To improve the efficiency and
performance of policy-based learning agents, various models are applied as
representations of the policy. Scalability, generality and differentiability
are desired characteristics of these models.

More importantly, the models with broader capability of representation are
essential for complex tasks. Among those policy representations, stochastic
policies are those formulated as a‘ probabilistic distribution, that can
generally be characterised as unimodal ones and multimodal ones. In prior
works, the choice of stochastic policies in DRL is typically with the unimodal
Gaussian distribution (Rawlik et al., 2013; Haarnoja et al., 2018; Schulman et
al., 2017). However, most complex tasks usually allow multiple optimal
solutions for even the same state, and a standard unimodal policy does not
have the capability of capturing or leveraging that. Thus, some methods (Peng
et al., 2019; Neumann et al., ; Peng et al., 2016; Akrour et al., 2020;
Hausknecht & Stone, 2016; Jacobs et al., 1991a) are proposed recently to model
the policy as a multimodal distribution with the mixture-of-experts (MOE). The
MOE, as originally proposed by (Jacobs et al., 1991b), can be implemented with
a Gaussian mixture model (GMM) in practice, which is also applicable for
representing the DRL policy. Each unimodal Gaussian component of GMM can be
seen as an expert or the so-called primitive (Sutton et al., 1999a), and the
mixing coefficients of GMM are the probability for activating each primitive
given an observation (Peng et al., 2019). The distinguishable primitives
learned by MOE can propose several solutions for a task, which may potentially
lead to better task performance and sample efficiency compared to its unimodal
counterpart.

However, a straightforward usage of GMM in generic off-policy and on-policy
DRL algorithms will face the indifferentiability problem in its end-to-end
training scheme, since optimising the GMM contains a process of optimising the
categorical distribution parameters (Bishop, 2007). Prior MOE methods for RL
usually works under some special settings, such as a two-stage training manner
(Peng et al., 2019), or with special task design (Hausknecht & Stone, 2016).
The indifferentiability problem can to solved to some extent by employing
their policy representations, but our experiments show that it generally does
not lead to better learning performances. One possible explanation is the
biased gradient estimation for the categorical distribution parameters of GMM.
Previous methods for solving the indifferentiability involves the
reparameterisation trick, _e.g._ , Gumbel-Softmax (Jang et al., 2017; Maddison
et al., 2017) and the score-ratio trick, _e.g._ , REINFORCE (Williams, 1992;
Glynn, 1990; Mnih & Rezende, 2016; Mnih & Gregor, 2014; Gu et al., 2015;
Gregor et al., 2014). However, the Gumbel-Softmax method will bring
brittleness to the temperature hyperparameter (Geffner & Domke, 2020; Tucker
et al., 2017) in the softmax function, while the score-ratio tricks are hard
to converge (Rezende et al., 2014b) due to the large variances. These problems
will be fatal for the fragile DRL training (Haarnoja et al., 2018) in general.

To tackle the above challenges, we propose a probabilistic MOE as a multimodal
function approximator, which can be used as DRL policies for generic off-
policy and on-policy DRL algorithms using stochastic policies. Moreover, a
novel gradient estimator named Frequency Approximate Gradient is proposed for
solving the indifferentiability problem when optimizing the MOE. Our method
has shown advantageous performances on 6 MuJoCo continuous control tasks with
an improvement up to 28.4% in the SAC-based experiments and 39.3% in the PPO-
based experiments compared to vanilla SAC and PPO respectively, measured with
the area under the curve (AUC). Experiments also show the advantageous
performance of our method over two other gradient estimation methods, Gumbel-
Softmax and REINFORCE. We further analyse the learned MOE policies by
displaying its distinguishable primitives and measuring its sensitivity to the
choice of primitive numbers, as well as giving a potential reason for the
performance improvement in terms of the exploration ranges.

##  2 Related Work

#### Mixture-of-Experts

To speed up the learning and improve the generalisation ability on different
scenarios, Jacobs et al. (1991a) proposed to use several different expert
networks instead of a single one. To partition the data space and assign
different kernels for different spaces, (Lima et al., 2007; Yao et al., 2009)
combine MOE with SVM. To break the dependency among training outputs and speed
up the convergence, the Gaussian process (GP) is generalised similarly to MOE
(Tresp, 2000; Yuan & Neubauer, 2008; Luo & Sun, 2017). MOE can be also
combined with RL (Doya et al., 2002; Neumann et al., ; Peng et al., 2016;
Hausknecht & Stone, 2016; Peng et al., 2019), in which the policies are
modelled as probabilistic mixture models and each expert aim to learn
distinguishable policies. (Peng et al., 2016) introduces a mixture of actor-
critic experts approaches to learn terrain-adaptive dynamic locomotion skills.
(Peng et al., 2019) changes the mixture-of-experts distribution addition
expression into the multiplication expression.

#### Hierarchical Policies

The MOE policy structure can be seen as a hierarchical structure, as the agent
chooses a Gaussian component of GMM to act according to the mixing
coefficient. There are two main related hierarchical policy structures. One is
the feudal schema (Dayan & Hinton, 1992), which has “manager” agents to first
make high-level decisions and the “worker” agents to make low-level actions.
(Vezhnevets et al., 2017) generalises the feudal schema into continuous action
space and uses an embedding operation to solve the indifferentiable problem.
The other is the options framework (Sutton et al., 1999b; McGovern & Barto,
2001; Menache et al., 2002; Simsek & Barto, 2008; Silver & Ciosek, 2012),
which has an upper-level agent (policy-over-options) to decide whether the
lower-level agent (sub-policy) should start or terminate. (Kulkarni et al.,
2016) uses internal and extrinsic rewards to learn sub-policies and policy-
over-options. (Bacon et al., 2017) trains sub-policies and policy-over-options
with a deep termination function.

#### Gradient Estimation Methods

Stochastic neural networks rarely use discrete latent variables due to the
inability to backpropagate through samples (Jang et al., 2017). Existing
stochastic gradient estimation methods traditionally focus on the Path
derivative gradient estimators and the Score-ratio gradient estimators. Path
derivative gradient estimators are formulated specifically for
reparameterisable distributions, _e.g._ , (Kingma & Welling, 2013; Rezende et
al., 2014a) employs a reparameterisation trick for the latent Gaussian
distribution, (Bengio et al., 2013) designs a Straight-Through estimator for
Bernoulli distribution, and (Jang et al., 2017) introduces Gumbel-Softmax to
approximate categorical samples whose parameter gradients can be easily
computed via the reparameterisation trick. Score-ratio gradient estimators use
the Log-derivative trick to derive an estimator, such as the score function
estimator(also referred to as REINFORCE (Williams, 1992)), likelihood ratio
estimator (Glynn, 1990), and other estimators augmented with Monte Carlo
variance reduction techniques (Mnih & Rezende, 2016; Mnih & Gregor, 2014; Gu
et al., 2015; Gregor et al., 2014).

#### Policy-based RL

Policy-based RL aims to find the optimal policy to maximise the expected
return through gradient updates. Among various algorithms, Actor-critic is
often employed (Barto et al., 1983; Sutton & Barto, 1998). Off-policy
algorithms (O’Donoghue et al., 2016; Lillicrap et al., 2016; Gu et al., 2017;
Haarnoja et al., 2018) are more sample efficient than on-policy ones (Peters &
Schaal, 2008; Schulman et al., 2017; Mnih et al., 2016; Gruslys et al., 2017).
However, the learned policies are still unimodal.

##  3 Method

###  3.1 Notation

The model-free RL problem can be formulated by Markov Decision Process (MDP),
denoted as a tuple $(\mathcal{S},\mathcal{A},P,r)$, where $\mathcal{S}$ and
$\mathcal{A}$ are continuous state and action space, respectively. The agent
observes state $s_{t}\in\mathcal{S}$ and takes an action $a_{t}\in\mathcal{A}$
at time step $t$. The environment emits a reward
$r:\mathcal{S}\times\mathcal{A}\rightarrow[r_{min},r_{max}]$ and transitions
to a new state $s_{t+1}$ according to the transition probabilities
$P:\mathcal{S}\times\mathcal{S}\times\mathcal{A}\rightarrow[0,\infty)$. In
deep reinforcement learning algorithms, we always use the Q-value function
$Q(s_{t},a_{t})$ to describe the expected return after taking an action
$a_{t}$ in the state $s_{t}$. The Q-value can be iteratively computed by
applying the Bellman backup given by:

$Q(s_{t},a_{t})\triangleq\mathbb{E}_{s_{t+1}\sim
P}\big{[}r(s_{t},a_{t})+\gamma\mathbb{E}_{a_{t+1}\sim\pi}[Q(s_{t+1},a_{t+1})]\big{]}.$
(1)

Our goal is to maximise the expected return:

$\pi_{\Theta*}(a_{t}|s_{t})=\arg\max_{\pi_{\Theta}(a_{t}|s_{t})}\mathbb{E}_{a_{t}\sim\pi_{\Theta}(a_{t}|s_{t})}[Q(s_{t},a_{t})],$
(2)

where $\Theta$ denotes the parameters of the policy network $\pi$. With
Q-value network (critic) $Q_{\phi}$ parameterised by $\phi$, Stochastic
gradient descent (SGD) based approaches are usually used to update the policy
network:

$\Theta\leftarrow\Theta+\nabla_{\Theta}\mathbb{E}_{a\sim\pi_{\Theta}(a_{t}|s_{t})}[Q_{\phi}(s_{t},a_{t})].$
(3)

###  3.2 Probabilistic Mixture-of-Experts (PMOE)

The proposed PMOE method follows the typical setting of MOE method, which
decomposes a complex policy $\pi$ into a mixture of low-level stochastic
policies with each of them as a probability distribution, represented as the
following:

$\displaystyle\pi_{\\{\theta,\psi\\}}(a_{t}|s_{t})$
$\displaystyle=\sum_{i=1}^{K}w_{\theta_{i}}(s_{t})\pi_{\psi_{i}}(a_{t}|s_{t}),$
(4) $\displaystyle s.t.\,\sum_{i=1}^{K}w_{\theta_{i}}=1,\,w_{\theta_{i}}>0,$
(5)

where each $\pi_{\psi_{i}}$ denotes the action distribution within each low-
level policy, _i.e._ a primitive, and $K$ denotes the number of primitives.
$w_{\theta_{i}}$ is the weight that specifies the probability of the
activating primitive $\pi_{\psi_{i}}$, which is called the routing function.
According to the GMM assumption (Bishop, 2007), $w_{\theta}$ is a Categorical
distribution and $\pi_{\psi}$ is a unimodal Gaussian distribution. For
$\forall i\in\\{1,2,...,K\\}$, $\theta_{i}$ and $\psi_{i}$ are parameters of
$w_{\theta_{i}}$ and $\pi_{\psi_{i}}$, respectively. After the policy
decomposition with PMOE method, we can rewrite the update rule in Eq. 3 as:

$\displaystyle\theta\leftarrow\theta+\nabla_{\theta}\mathbb{E}_{a_{t}\sim\pi_{\\{\theta,\psi\\}}(a_{t}|s_{t})}[Q_{\phi}(s_{t},a_{t})],$
(6)
$\displaystyle\psi\leftarrow\psi+\nabla_{\psi}\mathbb{E}_{a_{t}\sim\pi_{\\{\theta,\psi\\}}(a_{t}|s_{t})}[Q_{\phi}(s_{t},a_{t})].$

In practice, if we apply a Gaussian distribution for each of the low-level
policies here in PMOE, the overall PMOE will end up to be a GMM. However,
sampling from the mixture distributions of primitives
$\pi_{\\{\theta,\psi\\}}(a_{t}|s_{t})$ embeds a sampling process from the
categorical distribution $w_{\theta}$, which makes the differential
calculation of policy gradients commonly applied in DRL hard to achieve. We
provide a theoretically guaranteed solution for approximating the gradients in
the sampling process of PMOE and apply it for optimising the PMOE policy model
within DRL, which will be described in details.

###  3.3 Learning the Routing

The routing function in MOE typically involves a sampling process from a
categorical distribution (due to the discontinuity among multiple experts),
which is indifferentiable (Jang et al., 2017). To handle this difficulty, we
propose a new gradient estimator for this routing function.

Specifically, given a state $s_{t}$, we sample one action $a_{t}^{i}$ from
each primitive $\pi_{\psi_{i}}$, to get a total of $K$ actions
$\\{a_{t}^{i};i=1,2,\cdots,K\\}$, and compute $K$ Q-value estimations
$\\{Q_{\phi}(s_{t},a_{t}^{i});i=1,2,\cdots,K\\}$ for each of the actions. We
say the primitive $j$ is the “optimal” primitive under the Q-value estimation
if $j=\arg\max_{i}Q_{\phi}(s_{t},a_{t}^{i})$. There exists a frequency of the
primitive $j$ to be the “optimal” primitive given a set of state
$\\{s_{t}\\}$, here we propose a new gradient estimator which optimise
$\theta_{k}$ towards the frequency.

######  Definition 3.1 (Frequency Approximate Gradient).

For a stochastic mixture-of-experts, the gradient value of a single-instance
sampling process for the routing function $w_{\theta}$ can be estimated with
the frequency approximate gradient, which is defined as:

$\text{grad}=\delta_{k}\nabla_{\theta_{k}}w_{\theta_{k}},\,\delta_{k}=-\mathbbm{1}_{k}^{\text{best}}+w_{\theta_{k}},$
(7)

where $\nabla_{\theta_{k}}w_{\theta_{k}}$ is the gradient of $w_{\theta_{k}}$
for parameters $\theta_{k}$ and $\mathbbm{1}_{k}^{\text{best}}$ is the
indicator function that $\mathbbm{1}_{k}^{\text{best}}=1$ if
$k=\arg\max_{j}Q_{\phi}(s_{t},a_{t}^{j})$ and
$\mathbbm{1}_{k}^{\text{best}}=0$ otherwise.

######  Theorem 3.1.

The accumulated frequency approximate gradient is an asymptotically unbiased
estimation of the true gradient for the sampling process from a categorical
distribution in the routing function, with a batch of $N\to\infty$ samples.

###### Proof.

$\displaystyle\delta_{k}$
$\displaystyle=-\mathbbm{1}_{k}^{\text{best}}+w_{\theta_{k}}$ (8)
$\displaystyle=\mathbbm{1}_{k}^{\text{best}}w_{\theta_{k}}-\mathbbm{1}_{k}^{\text{best}}w_{\theta_{k}}-\mathbbm{1}_{k}^{\text{best}}+w_{\theta_{k}}$
$\displaystyle=\frac{1}{2}\mathbbm{1}_{k}^{\text{best}}\nabla_{w_{\theta_{k}}}(1-w_{\theta_{k}})^{2}+\frac{1}{2}(1-\mathbbm{1}_{k}^{\text{best}})\nabla_{w_{\theta_{k}}}w_{\theta_{k}}^{2}.$

Now we assume a batch of samples with a number of $N$ is applied, and the true
probability of primitive $k$ to be the best primitive (_i.e._
$k=\arg\max_{j}Q_{\phi}(s_{t},a_{t}^{k})$, is $f_{k}$. The batch accumulated
gradient will be

$\displaystyle\overline{\text{grad}}=$
$\displaystyle\frac{1}{N}\sum_{j=1}^{N}\text{grad}$ (9) $\displaystyle=$
$\displaystyle\frac{N_{t}}{2N}\nabla_{w_{\theta_{k}}}(1-w_{\theta_{k}})^{2}\nabla_{\theta_{k}}w_{\theta_{k}}$
$\displaystyle+\frac{(N-N_{t})}{2N}\nabla_{w_{\theta_{k}}}w_{\theta_{k}}^{2}\nabla_{\theta_{k}}w_{\theta_{k}}$
$\displaystyle\stackrel{{\scriptstyle N\to\infty}}{{=}}$
$\displaystyle(f_{k}-w_{\theta_{k}})\nabla_{\theta_{k}}w_{\theta_{k}},$

where the last formula indicates that $N_{t}=Nf_{k}$ when $N\to\infty$, since
the true probability can be approximated by $\frac{N_{t}}{N}$ in the limit
case. Since $\nabla_{\theta_{k}}w_{\theta_{k}}$ is not always equal to $0$,
the gradient equals to $0$ if and only if $w_{\theta_{k}}=f_{k}$ . Optimising
with the above equation is the same as minimising the distance between
$w_{\theta_{k}}$ and $f_{k}$, with the optimal situation as
$w_{\theta_{k}}=f_{k}$ when letting the last formula of Eq. 9 be zero. ∎

Then $\theta_{k}$ is updated via a gradient descent based algorithm, _e.g._ ,
Stochastic Gradient Descent (SGD):

$\theta_{k}\leftarrow\theta_{k}-\delta_{k}\nabla_{\theta_{k}}w_{\theta_{k}}.$
(10)

According to the formation of Eq 8, we can build an elegant loss function that
achieves the same goal of gradient estimating:

$\mathcal{L}_{freq}=(v-w)(v-w)^{T},w=[w_{\theta_{1}},w_{\theta_{2}},\cdots,w_{\theta_{K}}],$
(11)

$v$ is a one-hot code vector $v=[v_{1},v_{2},\cdots,v_{K}]$ with:

$v_{j}=\left\\{\begin{aligned}
&1,\,\text{if}\,j=\arg\max_{i}Q_{\phi}(s_{t},a_{t}^{i});\\\
&0,\,\text{otherwise}.\end{aligned}\right.$ (12)

###  3.4 Learning the Primitive

To update the $\psi_{i}$ within each primitive, we provide two approaches of
optimising the primitives: back-propagation-all and back-propagation-max
manners.

For the back-propagation-all approach, we update all the primitive:

$\mathcal{L}_{pri}^{bpa}=-\sum_{i}^{K}Q_{\phi}(s_{t},a_{t}^{i}),\,a_{t}^{i}\sim\pi_{\psi_{i}}(a_{t}|s_{t}).$
(13)

For the back-propagation-max approach, we use the highest Q-value estimation
as the primitive loss:

$\mathcal{L}_{pri}^{bpm}=-\max_{i=1,2,\cdots,K}\\{Q_{\phi}(s_{t},a_{t}^{i})\\},\,a_{t}^{i}\sim\pi_{\psi_{i}}(a_{t}|s_{t}).$
(14)

With either approach, we have the same stochastic policy gradients as
following:

$\displaystyle\nabla_{\psi_{i}}\mathcal{L}_{pri}=$
$\displaystyle-\nabla_{\psi_{j}}\mathbb{E}_{\pi_{\psi_{i}}}[Q_{\phi}(s_{t},a_{t})]$
(15) $\displaystyle=$
$\displaystyle\mathbb{E}_{\pi_{\psi_{i}}}[-Q_{\phi}(s_{t},a_{t})\nabla_{\psi_{j}}\log\pi_{\psi_{j}}(a_{t}|s_{t})]$

Ideally, both approaches are feasible for learning a PMOE model. However, in
practice, we find that the back-propagation-all approach will tend to learn
primitives that are close to each other, while the back-propagation-max
approach is capable of keeping primitives distinguishable. The phenomenon is
demonstrated in our experimental analysis. Therefore, we adopt the back-
propagation-max approach as the default setting of PMOE model without
additional clarification.

###  3.5 Learning the Critic

Similar to the standard off-policy RL algorithms, our Q-value network is also
trained to minimise the Bellman residual:

$\displaystyle\mathcal{L}_{critic}=$
$\displaystyle\mathbb{E}_{(s_{t},a_{t})\sim\mathcal{D}}[\|Q_{\phi}(s_{t},a_{t})-[r_{t}+$
(16)
$\displaystyle\gamma\max_{a_{t+1}}Q_{\bar{\phi}}(s_{t+1},a_{t+1})]\|_{2}],a_{t+1}\sim\pi(a_{t+1}|s_{t+1})$

where $\bar{\phi}$ is the parameters of the target network.

The learning component can be easily embedded into the popular actor-critic
algorithms, such as soft actor-critic (SAC) (Haarnoja et al., 2018), one of
the state-of-the-art off-policy RL algorithms. In SAC,
$Q_{\psi}(s_{t},a_{t}^{j})$ is substituted with
$Q_{\psi}(s_{t},a_{t}^{j})+\alpha\mathcal{H}_{j}$, where $\alpha$ is
temperature and $\mathcal{H}_{j}=-\log\pi_{\psi_{j}}(a_{t}|s_{t})$ is the
entropy which are the same as in SAC. The algorithm is summarised in Algorithm
1. When $K=1$, our algorithm simply reduces to the standard SAC.

Algorithm 1 Probabilistic MOE

Input: $\theta,\psi,\phi$

Initialise the target networks
$\bar{\theta}\leftarrow\theta,\bar{\psi}\leftarrow\psi,\bar{\phi}\leftarrow\phi$

Initialise an empty replay buffer: $\mathcal{D}\leftarrow\Phi$

repeat

for each environment step do

Sample action from policy: $a_{t}\sim\pi_{\\{\theta,\psi\\}}(a_{t}|s_{t})$

Interact with the environment: $s_{t+1}=p(s_{t+1|a_{t},s_{t}})$

Store in buffer:
$\mathcal{D}=\mathcal{D}\cup\\{s_{t},a_{t},s_{t+1},r(s_{t},a_{t})\\}$

end for

for each update step do

Sample from buffer: $\\{s_{t},a_{t},s_{t+1},r_{t}\\}\sim\mathcal{D}$

Compute $\mathcal{L}_{freq}$ with Eq. 11, $\mathcal{L}_{pri}$ with Eq. 14, and
$\mathcal{L}_{critic}$ with Eq. 16

Update the policy with:
$\theta\leftarrow\theta-\lambda_{\theta}\nabla_{\theta}\mathcal{L}_{freq}$,
$\psi\leftarrow\psi-\lambda_{\psi}\nabla_{\psi}\mathcal{L}_{pri}$

Update the critic with:
$\phi\leftarrow\phi-\lambda_{\phi}\nabla_{\phi}\mathcal{L}_{critic}$

end for

until converge

Output: $\theta,\psi,\phi.$

##  4 Experiments

To testify the performance of our method, we conduct a thorough comparison on
a set of challenging continuous control tasks in OpenAI Gym MuJoCo
environments (Brockman et al., 2016) with other baselines, including a MOE
method with gating operation (Jacobs et al., 1991b), Double Option Actor-
Critic (DAC) (Zhang & Whiteson, 2019) option framework, the Multiplicative
Compositional Policies (MCP) (Peng et al., 2019), and other two
implementations of PMOE with Gumbel-Softmax (Jang et al., 2017; Maddison et
al., 2017) and REINFORCE (Williams, 1992). Well-known sample efficient
algorithms involving Soft Actor-Critic (SAC) (Haarnoja et al., 2018) and
Proximal Policy Optimisation (PPO) (Schulman et al., 2017) are adopted as
basic DRL algorithms in our experiments, where different multimodal policy
approximation methods are built on top. This verifies the generality of PMOE
for different DRL algorithms.

To experimentally show the properties learned with PMOE, a deeper
investigation of PMOE is conducted to find out the additional effects caused
by deploying mixture models rather than a single model in policy learning. We
start with a simple self-defined _target-reaching_ task with sparse rewards to
show our method can indeed learn various optimal solutions with
distinguishable primitives, which are further demonstrated on complicated
tasks in MuJoCo. Additionally, the exploration behaviours are also compared
between PMOE and other baselines to explain the advantageous learning
efficiency of PMOE. Finally, to know how the number of primitives affects the
performance, we test PMOE with different values of $K$ on the
HumanoidStandup-v2 environment to analyse the impact of the different number
of primitives.

###  4.1 Performance Evaluation

(a) Ant-v2 (b) Hopper-v2 (c) Walker2D-v2 (d) Humanoid-v2

(e) HumanoidStandup-v2 (f) HalfCheetah-v2

Figure 1: Training curves on MuJoCo benchmarks with SAC-based algorithms. We
set PMOE with $K=4$ in all the experiments except HalfCheetah-v2 with $K=2$
and HumanoidStandup-v2 and Humanoid-v2 with $K=10$.

The evaluation on the average returns with SAC-based and PPO-based algorithms
are shown in Fig. 1 and Fig. 2, respectively. Specifically, in Fig. 1, SAC-
based algorithms with either unimodal policy or our PMOE for policy
approximation are compared against the MCP (Peng et al., 2019) and gating
operation methods (Jacobs et al., 1991b) in terms of the average returns
across six typical MuJoCo tasks. In Fig. 2, a comparison with similar settings
while basing on a different DRL algorithm is conducted to show the generality
of PMOE method. Specifically, we compare the proposed PMOE for policy
approximation on PPO with DAC and PPO-based MCP methods . Besides, to
demonstrate that our novel gradient estimator is proper for PMOE, we also
compare our method with other two implementations based on the Gumbel-Softmax
and REINFORCE in Fig. 1 and in Fig. 2. PMOE is testified to provide
improvement for general DRL algorithms with stochastic policies on a variety
of tasks. Training details are provided in Appendix C.

(a) Ant-v2 (b) Hopper-v2 (c) Walker2D-v2 (d) Humanoid-v2

(e) HumanoidStandup-v2 (f) HalfCheetah-v2

Figure 2: Training curves on MuJoCo benchmark with PPO-based algorithms. We
set a larger number, $K=16$ for Ant-v2, $K=12$ for Hopper-v2, $K=4$ for
Walker2D-v2 and HumainoidStandup-v2, $K=8$ for Humanoid-v2 and $K=8$ for
HalfCheetah-v2.

#### Comparison of AUC

We compute the AUC (the area under the learning curve) to make the Fig. 1,
Fig. 2, Fig. 6 and Fig. 7 more readable. For SAC-based experiments, we assume
the AUC of SAC is 1, and AUC values for all methods are shown in Table 1. For
PPO-based experiments, we assume the AUC of PPO is 1, and AUC values for all
methods are shown in Table 2.

|  Walker2D |  Half-Cheetah |  Humanoid |  Humanoid-Standup |  Ant |  Hopper
---|---|---|---|---|---|---
SAC |  100.0% |  100.0% |  100.0% |  100.0% |  100.0% |  100.0%
Gating |  93.4% |  95.1% |  91.1% |  92.5% |  88.4% |  110.0%
MCP-SAC |  109.0% |  103.2% |  97.6% |  98.3% |  96.5% |  100.5%
REINFORCE |  60.0% |  60.7% |  9.8% |  73.0% |  71.5% |  63.7%
Gumbel |  73.3% |  79.7% |  115.7% |  96.7% |  63.9% |  111.9%
PMOE-SAC |  124.9% |  99.6% |  113.8% |  110.5% |  128.4% |  115.5%
Table 1: Comparison of the AUC between PMOE-SAC and other methods on six MuJoCo tasks. |  Walker2D |  Half-Cheetah |  Humanoid |  Humanoid-Standup |  Ant |  Hopper
---|---|---|---|---|---|---
PPO |  100.0% |  100.0% |  100.0% |  100.0% |  100.0% |  100.0%
DAC |  73.9% |  75.2% |  61.0% |  57.6% |  87.9% |  90.9%
MCP-PPO |  111.1% |  118.1% |  113.6% |  66.5% |  164.3% |  82.7%
REINFORCE |  63.0% |  103.6% |  102.3% |  89.1% |  124.6% |  90.5%
Gumbel |  88.9% |  83.4% |  100.2% |  103.8% |  81.2% |  44.9%
PMOE-PPO |  138.4% |  126.6% |  107.5% |  105.3% |  139.3% |  120.7%
Table 2: Comparison of the AUC between PMOE-PPO and other methods on six
MuJoCo tasks.

###  4.2 Experiment Analysis

We provide in-depth analysis of the proposed PMOE to analysis the differences
between PMOE method and other methods in RL process, and estimate the effects
of the number of primitives.

#### Distinguishable Primitives.

Firstly, for a straight illustration of the distinguishable primitives, we use
a self-defined _target-reaching_ sparse reward environment as a toy example to
analyse our method. In this environment, the agent starts from a fixed initial
position, then acts to reach a random target position within a certain range
and avoids the obstacles on the path. Only when the agent reaches the target
position successfully, the agent can get the reward. As the reward setting is
sparse, the exploration can be important, so we also analyse the exploration
behaviours in this environment. Details about this environment is provided in
Appendix B.

Fig. 3 demonstrates the distinguishable primitives learned with PMOE on the
target-reaching environment, for providing a simple and intuitive
understanding. After the training stage, we sample 10 trajectories for each
method and visualise them in Fig. 3. As we can see in Fig. 3(e), PMOE trained
in back-propagation-max manner generates two distinguishable trajectories for
different primitives.

In Fig. 4, we further demonstrate that PMOE can learn distinguishable
primitives on more complex MuJoCo environments with the t-SNE (van der Maaten
& Hinton, 2008) method for visualisation. We sample 10K states
$\\{s_{t};t=1,2,\cdots,10K\\}$ from 10 trajectories and use t-SNE to perform
dimensionality reduction on states and visualise the results in Fig. 4(b). We
randomly choose one state cluster and sample actions corresponding to the
states in that cluster. Then we use t-SNE to perform dimensionality reduction
on those actions. The reason for taking a state clustering process before the
action clustering is to reduce the variances of data in state-action spaces so
that we can better visualise the action primitives for a specific cluster of
states. The visualisation of action clustering with our approach and the
gating operation are displayed in Fig. 4(a) and Fig. 4(c). More t-SNE
visualisations for other MuJoCo environments can be found in Appendix E. Our
proposed PMOE method is testified to have stronger capability in
distinguishing different primitives during policy learning process.

(a) (b) (c) (d)

(e)

Figure 3: Trajectories of the agents with our method and the baselines in the
target-reaching environment. We fix the reset locations of target, obstacles
and agent. ($a$), ($b$), ($c$) and ($d$) visualise the 10 trajectories
collected with methods involving: original SAC, gating operation with SAC,
back-propagation-all PMOE (discussed in Sec. 14) and back-propagation-max
PMOE, respectively. ($e$) shows the trajectories collected with two individual
primitives with our approach.

(a) Gating operation.

(b) States clusters.

(c) Our approach.

Figure 4: Visualisation of distinguishable primitives learned with PMOE using
t-SNE plot on Hopper-v2 environment. The states are first clustered as in
($b$). Then actions within the same state cluster are plotted with t-SNE as in
($a$) and ($c$) for the gating method and our approach, respectively. Our
method clearly demonstrates more distinguishable primitives for the policy.

#### Exploration Behaviours.

Fig. 5 demonstrates the exploration trajectories in the target-reaching
environment, although all trajectories start from the same initial positions,
our methods demonstrate larger exploration ranges compared against other
baseline methods, which also yields a higher visiting frequency to the target
region (in green) and therefore accelerates the learning process. To some
extent, this comparison can be one reason for the improvement of using PMOE as
the policy representation over a general unimodal policy. We find that by
leveraging the mixture models, the agents gain effective information more
quickly via different exploration behaviours, which cover a larger range of
exploration space at the initial stages of learning and therefore ensure a
higher ratio of target reaching.

(a) Ours (b) Success Rate in Learning (c) SAC

Figure 5: Visualisation of exploration trajectories in the initial training
stage for the target-reaching environment. The initial $10K$ steps (the grey
region on the learning curves in ($b$)) of exploration trajectories are
plotted in ($a$) and ($c$) for our PMOE method (red) and SAC (blue),
respectively. The green rectangle is the target region. Figure 6:  Comparison
of different numbers of primitives $K$ in terms of average returns on
HumanoidStandup-v2 environment. For each case, we conduct 5 runs and take the
means. The performance increases when $K$ increase from $2$ to $10$, but
decreases if $K$ keep increasing from $10$ to $16$.

#### Number of Primitives.

We investigate the effects caused by different numbers of primitives in GMM,
as shown in Fig. 6. This experiment is conducted on a relatively complex
environment — HumanoidStandup-v2 that has observations with a dimension of 376
and actions with a dimension of 17, therefore various skills could possibly
lead to the goal of the task. The number of primitives $K$ is selected from
$\\{2,4,8,10,14,16\\}$ and other hyperparameters are the same. The results
show that $K=10$ seems to perform the best, and $K=2$ performs the worst among
all settings, showing that increasing the number of primitives can improve the
learning efficiency in some situations. For Table 3, we assume the AUC of
PMOE-SAC with $K=2$ is 1, and relative AUC values for all the methods are
displayed.

Number of K | AUC
---|---
2 | 100%
4 | 107.1%
8 | 106.5%
10 |  111.7%
14 | 108.0%
16 | 107.0%
Table 3: Comparison of the AUC as a function of $K$ for PMOE-SAC algorithm on
HumanoidStandup-v2 environment.

To analyse the relationship of different $K$ and different entropy
regularisation($\alpha$), we also compared 7 settings in the MuJoCo task
HumanoidStandup-v2, where $K$ is the number of primitives and $\alpha$ is the
entropy regularisation. For each setting, we randomly choose 5 seeds to plot
the learning curves in Fig. 7. For Table 4, weassume the AUC of SAC with
$\alpha$ = 10 is 1, and relative AUC values for all the methods are displayed.

Settings | AUC
---|---
SAC, $\alpha$=10 | 100.0%
SAC, $\alpha$=1 | 121.7%
SAC, $\alpha$=0.2 | 108.3%
SAC, $\alpha$=0.05 | 108.5%
K=4, $\alpha$=10 | 57.3%
K=4, $\alpha$=1 | 125.7%
K=4, $\alpha$=0.2 | 114.7%
K=4, $\alpha$=0.05 | 116.0%
Table 4: Comparison of the AUC as a function of the relation between $K$ and
$\alpha$ for PMOE-SAC algorithm on HumanoidStandup-v2 environment. Figure 7:
Comparison with different $K$ and different amounts of entropy regularisation.
Our approach can be considered as a kind of entropy regularisation method and
the number of primitives is positively correlated with the entropy of the
policy. The larger number of primitives with smaller entropy has a similar
performance to the smaller number of primitives with larger entropy.

#### Robustness Evaluation

To evaluate the robustness of our approach, we develop an experiment on the
Hopper-v2 environment. We add a random noise
$\epsilon\sim\mathcal{N}(0,\sigma^{2})$ to the state observation $s$, and use
the noised state observation $\hat{s}=s+\epsilon$ as the input of the policy.
Our approach has a more stable performance in the noised input observation
situation, which is shown in Table 5.

Method | $\sigma=0$ | $\sigma=0.05$ | $\sigma=0.1$
---|---|---|---
SAC |  3387.4 $\pm$ 2.0 |  1994.9 $\pm$ 718.6 |  1006.2 $\pm$ 389.6
gating |  3444.8 $\pm$ 3.1 |  2606.9 $\pm$ 864.4 |  1626.0 $\pm$ 771.9
MCP |  3524.8 $\pm$ 114.6 |  1610.2 $\pm$ 357.0 |  1008.4 $\pm$ 333.4
Gumbel-Softmax |  3248.3$\pm$453.6 |  3042.9 $\pm$937.5 |  1347.4$\pm$855.0
REINFORECE |  1741.5$\pm$618.1 |  1111.6$\pm$398.1 |  558.7$\pm$276.5
Ours |  3632.2 $\pm$ 4.0 |  3460.4 $\pm$ 456.7 |  1730.0 $\pm$ 703.1
Table 5: We test our approach in the Hopper-v2 environment, each column stands
for the average return with different variances of the noise distribution, the
average return of each methods is averaged over 100 rounds.

##  5 Conclusion

To cope with the problems of low learning efficiency and multimodal solutions
in continuous control tasks when applying DRL, this paper proposes the
differentiable PMOE method that enables an end-to-end training scheme for
generic RL algorithms with stochastic policies. Our proposed method is
compatible with policy-gradient-based algorithms, like SAC and PPO.
Experiments show performance improvement across various tasks is achieved by
applying our PMOE method for policy approximation, as well as displaying
distinguishable primitives for multiple solutions.

## References

  * Akrour et al. (2020) Akrour, R., Tateo, D., and Peters, J.  Reinforcement learning from a mixture of interpretable experts.  _CoRR_ , abs/2006.05911, 2020.
  * Bacon et al. (2017) Bacon, P., Harb, J., and Precup, D.  The option-critic architecture.  In _AAAI_ , pp. 1726–1734, 2017.
  * Barto et al. (1983) Barto, A. G., Sutton, R. S., and Anderson, C. W.  Neuronlike adaptive elements that can solve difficult learning control problems.  _IEEE Trans. Syst. Man Cybern._ , 13(5):834–846, 1983.
  * Bengio et al. (2013) Bengio, Y., Léonard, N., and Courville, A.  Estimating or propagating gradients through stochastic neurons for conditional computation.  _arXiv preprint arXiv:1308.3432_ , 2013.
  * Bishop (2007) Bishop, C. M.  _Pattern recognition and machine learning, 5th Edition_.  Information science and statistics. Springer, 2007.
  * Brockman et al. (2016) Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba, W.  Openai gym.  _CoRR_ , abs/1606.01540, 2016.
  * Dayan & Hinton (1992) Dayan, P. and Hinton, G. E.  Feudal reinforcement learning.  In _NeurIPS_ , pp. 271–278, 1992.
  * Dong et al. (2020) Dong, H., Ding, Z., Zhang, S., and Chang.  _Deep Reinforcement Learning_.  Springer, 2020.
  * Doya et al. (2002) Doya, K., Samejima, K., Katagiri, K., and Kawato, M.  Multiple model-based reinforcement learning.  _Neural Comput._ , 14(6):1347–1369, 2002.
  * Geffner & Domke (2020) Geffner, T. and Domke, J.  Approximation based variance reduction for reparameterization gradients.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.
  * Glynn (1990) Glynn, P. W.  Likelihood ratio gradient estimation for stochastic systems.  _Communications of the ACM_ , 33(10):75–84, 1990\.
  * Gregor et al. (2014) Gregor, K., Danihelka, I., Mnih, A., Blundell, C., and Wierstra, D.  Deep autoregressive networks.  In _International Conference on Machine Learning_ , pp. 1242–1250. PMLR, 2014.
  * Gruslys et al. (2017) Gruslys, A., Azar, M. G., Bellemare, M. G., and Munos, R.  The reactor: A sample-efficient actor-critic architecture.  _CoRR_ , abs/1704.04651, 2017.
  * Gu et al. (2015) Gu, S., Levine, S., Sutskever, I., and Mnih, A.  Muprop: Unbiased backpropagation for stochastic neural networks.  _arXiv preprint arXiv:1511.05176_ , 2015.
  * Gu et al. (2017) Gu, S., Lillicrap, T. P., Ghahramani, Z., Turner, R. E., and Levine, S.  Q-prop: Sample-efficient policy gradient with an off-policy critic.  In _ICLR_ , 2017.
  * Haarnoja et al. (2018) Haarnoja, T., Aurick, Z., Pieter, A., and Sergey, L.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  volume 80, pp. 1861–1870. PMLR, 10–15 Jul 2018.
  * Hausknecht & Stone (2016) Hausknecht, M. J. and Stone, P.  Deep reinforcement learning in parameterized action space.  In _ICLR_ , 2016.
  * Jacobs et al. (1991a) Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E.  Adaptive mixtures of local experts.  _Neural Comput._ , 3(1):79–87, 1991a.
  * Jacobs et al. (1991b) Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E.  Adaptive mixtures of local experts.  _Neural Comput._ , 3(1):79–87, 1991b.
  * Jang et al. (2017) Jang, E., Gu, S., and Poole, B.  Categorical reparameterization with gumbel-softmax.  In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net, 2017.
  * Kingma & Ba (2015) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  In _ICLR_ , 2015.
  * Kingma & Welling (2013) Kingma, D. P. and Welling, M.  Auto-encoding variational bayes, 2013.
  * Kulkarni et al. (2016) Kulkarni, T. D., Narasimhan, K., Saeedi, A., and Tenenbaum, J.  Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation.  In _NeurIPS_ , pp. 3675–3683, 2016.
  * Lillicrap et al. (2016) Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., and Wierstra, D.  Continuous control with deep reinforcement learning.  In _ICLR_ , 2016.
  * Lima et al. (2007) Lima, C. A. M., Coelho, A. L. V., and Zuben, F. J. V.  Hybridizing mixtures of experts with support vector machines: Investigation into nonlinear dynamic systems identification.  _Inf. Sci._ , 177(10):2049–2074, 2007.
  * Luo & Sun (2017) Luo, C. and Sun, S.  Variational mixtures of gaussian processes for classification.  In _IJCAI_ , pp. 4603–4609, 2017.
  * Maddison et al. (2017) Maddison, C. J., Mnih, A., and Teh, Y. W.  The concrete distribution: A continuous relaxation of discrete random variables.  In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net, 2017.
  * McGovern & Barto (2001) McGovern, A. and Barto, A. G.  Automatic discovery of subgoals in reinforcement learning using diverse density.  In _ICML_ , pp. 361–368, 2001.
  * Menache et al. (2002) Menache, I., Mannor, S., and Shimkin, N.  Q-cut - dynamic discovery of sub-goals in reinforcement learning.  In _ECML_ , volume 2430, pp. 295–306, 2002.
  * Mnih & Gregor (2014) Mnih, A. and Gregor, K.  Neural variational inference and learning in belief networks.  In _International Conference on Machine Learning_ , pp. 1791–1799. PMLR, 2014.
  * Mnih & Rezende (2016) Mnih, A. and Rezende, D.  Variational inference for monte carlo objectives.  In _International Conference on Machine Learning_ , pp. 2188–2196. PMLR, 2016.
  * Mnih et al. (2016) Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T. P., Harley, T., Silver, D., and Kavukcuoglu, K.  Asynchronous methods for deep reinforcement learning.  In _ICML_ , volume 48, pp. 1928–1937, 2016.
  * (33) Neumann, G., Maass, W., and Peters, J.  Learning complex motions by sequencing simpler motion templates.  In _ICML_ , volume 382, pp. 753–760.
  * O’Donoghue et al. (2016) O’Donoghue, B., Munos, R., Kavukcuoglu, K., and Mnih, V.  PGQ: combining policy gradient and q-learning.  _CoRR_ , abs/1611.01626, 2016.
  * Peng et al. (2016) Peng, X. B., Berseth, G., and van de Panne, M.  Terrain-adaptive locomotion skills using deep reinforcement learning.  _ACM Trans. Graph._ , 35(4):81:1–81:12, 2016\.
  * Peng et al. (2019) Peng, X. B., Chang, M., Zhang, G., Abbeel, P., and Levine, S.  MCP: learning composable hierarchical control with multiplicative compositional policies.  In _NeurIPS_ , pp. 3681–3692, 2019.
  * Peters & Schaal (2008) Peters, J. and Schaal, S.  Reinforcement learning of motor skills with policy gradients.  _Neural Networks_ , 21(4):682–697, 2008.
  * Rawlik et al. (2013) Rawlik, K., Toussaint, M., and Vijayakumar, S.  On stochastic optimal control and reinforcement learning by approximate inference (extended abstract).  In Rossi, F. (ed.), _IJCAI 2013, Proceedings of the 23rd International Joint Conference on Artificial Intelligence, Beijing, China, August 3-9, 2013_ , pp. 3052–3056. IJCAI/AAAI, 2013.
  * Rezende et al. (2014a) Rezende, D. J., Mohamed, S., and Wierstra, D.  Stochastic backpropagation and approximate inference in deep generative models.  In _International conference on machine learning_ , pp. 1278–1286. PMLR, 2014a.
  * Rezende et al. (2014b) Rezende, D. J., Mohamed, S., and Wierstra, D.  Stochastic backpropagation and approximate inference in deep generative models.  In _International conference on machine learning_ , pp. 1278–1286. PMLR, 2014b.
  * Schulman et al. (2017) Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O.  Proximal policy optimization algorithms.  _CoRR_ , abs/1707.06347, 2017.
  * Silver & Ciosek (2012) Silver, D. and Ciosek, K.  Compositional planning using optimal option models.  In _ICML_ , 2012.
  * Simsek & Barto (2008) Simsek, Ö. and Barto, A. G.  Skill characterization based on betweenness.  In _NeurIPS_ , pp. 1497–1504, 2008.
  * Sutton & Barto (1998) Sutton, R. S. and Barto, A. G.  _Reinforcement learning - an introduction_.  Adaptive computation and machine learning. MIT Press, 1998.
  * Sutton et al. (1999a) Sutton, R. S., Precup, D., and Singh, S.  Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning.  _Artificial intelligence_ , 112(1-2):181–211, 1999a.
  * Sutton et al. (1999b) Sutton, R. S., Precup, D., and Singh, S. P.  Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning.  _Artif. Intell._ , 112(1-2):181–211, 1999b.
  * Tresp (2000) Tresp, V.  Mixtures of gaussian processes.  In _NeurIPS_ , pp. 654–660, 2000.
  * Tucker et al. (2017) Tucker, G., Mnih, A., Maddison, C. J., Lawson, D., and Sohl-Dickstein, J.  REBAR: low-variance, unbiased gradient estimates for discrete latent variable models.  In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pp. 2627–2636, 2017.
  * van der Maaten & Hinton (2008) van der Maaten, L. and Hinton, G.  Viualizing data using t-sne.  _JMLR_ , 9:2579–2605, 11 2008.
  * van Hasselt et al. (2015) van Hasselt, H., Guez, A., and Silver, D.  Deep reinforcement learning with double q-learning.  _CoRR_ , abs/1509.06461, 2015.
  * Vezhnevets et al. (2017) Vezhnevets, A. S., Osindero, S., Schaul, T., Heess, N., Jaderberg, M., Silver, D., and Kavukcuoglu, K.  Feudal networks for hierarchical reinforcement learning.  In _ICML_ , volume 70, pp. 3540–3549, 2017.
  * Wang et al. (2015) Wang, Z., de Freitas, N., and Lanctot, M.  Dueling network architectures for deep reinforcement learning.  _CoRR_ , abs/1511.06581, 2015.
  * Williams (1992) Williams, R. J.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine learning_ , 8(3-4):229–256, 1992.
  * Yao et al. (2009) Yao, B., Walther, D. B., Beck, D. M., and Li, F.  Hierarchical mixture of classification experts uncovers interactions between brain regions.  In _NeurIPS_ , pp. 2178–2186, 2009.
  * Yuan & Neubauer (2008) Yuan, C. and Neubauer, C.  Variational mixture of gaussian process experts.  In _NeurIPS_ , pp. 1897–1904, 2008.
  * Zhang & Whiteson (2019) Zhang, S. and Whiteson, S.  DAC: the double actor-critic architecture for learning options.  In _NeurIPS_ , pp. 2010–2020, 2019.

##  Appendix A Probabilistic formulation of PMOE and Gating Operation

In this section, we show a detailed comparison of probabilistic formulation
for GMM (as Eq. (17) and (18), used in our PMOE method) and the gating
operation method (Eq. (19) to (21)), in term of their PDFs. The gating
operation degenerates the multimodal action to a unimodal distribution, which
is different from our PMOE method.

For GMM, suppose a primitive $\pi_{i}(s)$ is a Gaussian distribution
$\mathcal{N}(a|\mu(s),\sigma^{2}(s))$, drawing a sample from the mixture model
can be seen as the following operation:

$\displaystyle a\sim\pi(a|s)$
$\displaystyle=\sum_{i=1}^{K}w_{i}(s)\pi_{i}(a|s)$ (17)
$\displaystyle=\sum_{i=1}^{K}w_{i}(s)\mathcal{N}(a|\mu_{i}(s),\sigma_{i}^{2}(s)),$

where the PDF is:

$p(a)=\sum_{i=1}^{K}\frac{w_{i}(s)}{\sqrt{2\pi}\sigma_{i}(s)}\exp\\{-\frac{(a-\mu_{i}(s))^{2}}{2\sigma_{i}^{2}(s)}\\}.$
(18)

For gating operation, the outputs of the weight operation are the weights of
each action from different primitives. With those weights, the gating
operation uses the weighted action as the final output action according to
(Peng et al., 2019):

$a=\sum_{i=1}^{K}w_{i}(s)a_{i},\,s.t.\,a_{i}\sim\pi_{i}(a|s)=\mathcal{N}(a|\mu_{i}(s),\sigma_{i}^{2}(s)).$
(19)

As a primitive is a Gaussian distribution, Eq. 19 becomes:

$a\sim\mathcal{N}(a|\sum_{i=1}^{K}w_{i}(s)\mu_{i}(s),\sum_{i=1}^{K}w_{i}(s)\sigma_{i}^{2}(s)),$
(20)

where the PDF is:

$\begin{split}p(a)=&\frac{1}{\sqrt{2\pi\sum_{i=1}^{K}w_{i}(s)\sigma_{i}^{2}(s)}}\\\
&\exp\\{-\frac{(a-\sum_{i=1}^{K}w_{i}(s)\mu_{i}(s))^{2}}{2\sum_{i=1}^{K}w_{i}(s)\sigma_{i}^{2}(s)}\\},\end{split}$
(21)

The above PDF shows that the gating operation could degenerate the Gaussian
mixture model into the univariate Gaussian distribution. Other methods (Jacobs
et al., 1991a; Peng et al., 2016; Vezhnevets et al., 2017) also have the
similar formulation.

##  Appendix B Details of Target-Reaching Environment

The visualisation of the target-reaching environment is shown in Fig.8, the
blue circle is the agent, the gray circles are obstacles and the circle in red
is the target. The agent state is represented by an action vector
$a=[a_{x},a_{y}]$ and a velocity vector $v=[v_{x},v_{y}]$. The playground of
the environment is continuous and limited in $[-5,5]$ in both x-axis and
y-axis. The agent speed is limited into $[-2,2]$, the blue coloured agent is
placed at position $[x_{ag},y_{ag}]=[-4.5,-4.5]$ and the red coloured target
is randomly placed at position $[x_{tg},y_{tg}]$, where
$x_{tg},y_{tg}\sim\mathcal{U}(0,3)$ and $\mathcal{U}$ denotes uniform
distribution. There are $M$ gray coloured obstacles with each position
$[x_{obs}^{i},y_{obs}^{i}]$, where
$x_{obs}^{i},y_{obs}^{i}\sim\mathcal{N}(0,3^{2})$ and $\mathcal{N}$ denotes
Gaussian distribution. The observation is composed of
$[[x_{tg}-x_{ag},y_{tg}-y_{ag}],\\{[x_{obs}^{i}-x_{ag},y_{obs}^{i}-y_{ag}^{i}];i=1,2,\cdots,M\\},a,v]$.
The input action is the continuous acceleration $a$ which is in the range of
$[-2,2]$.

Figure 8: Visualisation of the target-reaching environment

The immediate reward function for each time step is defined as:

$r=\left\\{\begin{aligned} &100,\,\text{if the agent reaches the target};\\\
&-10,\,\text{if the agent collides with edges or obstacles};\\\
&||v||_{2},\text{otherwise}.\end{aligned}\right.$ (22)

##  Appendix C Training Details

For PMOE-SAC policy network, we use a two-layer fully-connected (FC) network
with 256 hidden units and ReLU activation in each layer. For primitive network
$\pi_{\psi}$, we use a two single-layer FC network, which outputs $\mu$ and
$\sigma$ for the Gaussian distribution. Both the output layers for $\mu$ and
$\sigma$ have the same number of units, which is $K*dim(\mathcal{A})$, with
$K$ as the number of primitives and $dim(\mathcal{A})$ as the dimension of
action space, _e.g._ , 17 for Humanoid-v2. For the routing function network
$w_{\theta}$, we use a single FC layer with $K$ hidden units and the softmax
activation function.. In critic network we use a three-layer FC network with
256, 256 and 1 hidden units in each layer and ReLU activation for the first
two layers. Other hyperparameters for training are showed in Tab. 1(a). For
PMOE-PPO, we use a two-layer FC network to extract the features of state
observations. The FC layers have 64 and 64 hidden units with ReLU activation.
The policy network has a single layer with the Tanh activation function. The
routing function network has a single FC layer with $K$ units and the softmax
activation function. The critic contains one layer only. Other training
hyperparameters are showed in Tab. 1(b). We use the same hyperparameters in
all the experiments without any fine-tuning. For MCP-SAC, we use the same
network structure as MCP-PPO, other training hyperparemeters are the same as
shown in Tab. 1(a). For other baselines, we use original hyperparameters
mentioned in their paper. The full algorithm is summarised in Algorithm 1.

Parameter | Value
---|---
optimiser | Adam (Kingma & Ba, 2015)
learning rate | $10^{-3}$
discount ($\gamma$) | 0.99
replay buffer size | $10^{6}$
alpha | 0.2
batch size | 100
polyak ($\tau$) | 0.995
episode length | $10^{3}$
target update interval | 1
(a) Hyperparameters for PMOE-SAC

Parameter | Value
---|---
optimiser | Adam (Kingma & Ba, 2015)
learning rate | $3*10^{-4}$
discount ($\gamma$) | 0.99
alpha | 0.2
batch size | 64
polyak ($\tau$) | 0.95
episode length | $2*10^{3}$
gradient clip | 0.2
optimisation epochs | 20
(b) Hyperparameters for PMOE-PPO

Table 6: Hyperparameters

##  Appendix D Probability Visualisation

We visualise the probabilities of each primitive over the time steps in the
MuJoCo HalfCheetah-v2 environment. As shown in Fig. 9, we found that the
probabilities are changed periodically. We also visualise the actions at the
selected 5 time steps in one period. As shown in Fig. 10, the primitives are
distinguishable enough to develop distinct specialisations.

Figure 9: Visualisation of the probabilities of each primitive over the time
steps in the MuJoCo HalfCheetah-v2 environment. The y-axis shows the
probabilities of different primitives.

(a) time step 1

(b) time step 3

(c) time step 5

(d) time step 6

(e) time step 7

Figure 10: Visualisation of the actions at the selected 5 time steps in one
period. The y-axis shows the probabilities of different primitives. This
result shows that the primitives develop distinct specialisations, with the
primitive 0 becomes the most active when the front leg touches the ground,
while the primitive 1 becomes the most active when the leg leaves the ground.

##  Appendix E t-SNE Visualisation

To demonstrate the distinguishable primitives, we plot the t-SNE visualisation
for other 5 MuJoCo environments in Fig. 11.

(a) Gating Operation (b) State Observations (c) Ours

Figure 11: We plot the t-SNE visualisation for other 5 MuJoCo environments:
Ant-v2, HumanoidStandup-v2, Humanoid-v2, HalfCheetah-v2 and Walker2D-v2.
Parameters and other details are the same as the setting mentioned in Sec 4.