# We Need Variations in Speech Synthesis:
Sub-center Modelling for Speaker Embeddings

Ismail Rasim Ulgen, , Carlos Busso , John H. L. Hansen
and Berrak Sisman Manuscript submitted on 4 Jul. This work is supported by NSF
CAREER award IIS-2338979.The authors are with the Electrical and Computer
Engineering Department, University of Texas at Dallas, Richardson, TX 75080
USA (e-mail<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>berrak.sisman@utdallas.edu).

###### Abstract

In speech synthesis, modeling of rich emotions and prosodic variations present
in human voice are crucial to synthesize natural speech. Although speaker
embeddings have been widely used in personalized speech synthesis as
conditioning inputs, they are designed to lose variation to optimize speaker
recognition accuracy. Thus, they are suboptimal for speech synthesis in terms
of modeling the rich variations at the output speech distribution. In this
work, we propose a novel speaker embedding network which utilizes multiple
class centers in the speaker classification training rather than a single
class center as traditional embeddings. The proposed approach introduces
variations in the speaker embedding while retaining the speaker recognition
performance since model does not have to map all of the utterances of a
speaker into a single class center. We apply our proposed embedding in voice
conversion task and show that our method provides better naturalness and
prosody in synthesized speech.

###### Index Terms:

speaker embedding, speech synthesis, voice conversion, intra-class variance

##  I Introduction

Speaker embeddings are originally designed for the task of speaker
recognition, where the aim is to recognize the speaker’s identity from speech.
They are vectors obtained from one of the last hidden layers of deep neural
networks trained to classify speakers on a relatively large amount of data
with many speakers [1, 2, 3]. Since they are advanced deep neural networks
trained with a large amount of data, speaker embeddings are generalizable and
capable of encoding speaker information even from very limited data. These
strengths make speaker embeddings a very popular tool for various speech
synthesis tasks, including text-to-speech (TTS) [4, 5] and voice conversion
(VC) [6, 7, 8, 9, 10, 11].

Zero-shot, multi-speaker speech synthesis methods [6, 4, 7, 5] are becoming
increasingly popular as they are capable of synthesizing speech from any
speaker with very limited data, allowing personalization and modification of
synthesized speech for various use cases. There are two main zero-shot, multi-
speaker speech synthesis tasks that differ in terms of their input modality:
TTS and VC. In TTS, the input text is mapped into speech, while in VC, the
input speech is mapped into a target speaker’s voice while preserving the
linguistic content. In the zero-shot, multi-speaker setting, deep speaker
embeddings are widely used as conditioning inputs in the mapping in order to
make the model learn the speaker identity knowledge and be able to synthesize
speech from speakers that are unseen in the training. In speech synthesis
tasks, the main aim is to synthesize the most natural speech, and being able
to model the variations in speech (e.g., prosody, emotion, style) that change
voice characteristics of a speaker, is crucial to attaining that aim.

Since speaker recognition aims to distinguish a given speaker’s utterances
from others, the focus is on minimizing intra-class variance and maximizing
inter-class variance [12, 13]. However, minimizing intra-class variance often
results in the loss of variability between utterances by the same speaker
within the speaker embeddings. It is shown that traditional speaker embedding
networks with classification objective provide a single class center for each
class. At the last layer, the model attempts to make embedding of an utterance
close to the corresponding class center [14]. We note that pushing every class
member to one single center might result in losing valuable sub-class
variation information such as emotional and prosodic cues which are essential
for naturalistic speech synthesis. The absence of variation within speaker
embeddings might lead to suboptimal speech synthesis as the synthesis method
has to model the variations in the output speech distribution of a speaker
with very limited knowledge at the input. ††Speech samples: https://lec-
synt.github.io/sub-center-demo/

Recently, classification with sub-class centers has been introduced [14, 15]
for fine-grained image retrieval where the variation and intra-class modeling
is also important for well-sorted class members. In this work, we claim that a
larger intra-class variance in speaker embeddings might be more suitable for
synthesis tasks, as in the case of fine-grained image retrieval [16, 17]. We
show that sub-centers allows introducing variation in speaker embeddings as
all of utterances of a given speaker does not need to be collapsed into a
single-class center.

In this work, we apply sub-class center modeling to state-of-the-art speaker
embeddings for better intra-class variance modeling for zero-shot, multi-
speaker speech synthesis tasks. We show that sub-center modeling allows for
higher intra-class variance without compromision recognition performance; in
fact, it leads to improvement. We apply our speaker embeddings in the voice
conversion task, show improvement on synthesized speech. Our contributions can
be summarized as follows: 1) We introduce novel speaker embeddings with
subcenter modeling; 2) We provide a novel analysis of speaker embeddings in
speech synthesis from the intra-class variance perspective; and 3) We show
that when applied to voice conversion, the proposed embeddings with higher
intra-class variance exhibit more naturalness and consistent prosody compared
to embeddings with less intra-class variance.

##  II Related Work

Sub-center classification is relatively more popular in computer vision.
Studies such as [14] and [15] have utilized sub-center modeling and improved
fine-grained image retrieval. In [18], the authors applied sub-center modeling
to the teacher-student modeling problem, demonstrating that the teacher model
can learn sub-classes despite being trained only with parent class labels.
Diverging from these approaches, the authors of [19] focused on separating the
noisy labeled class examples from correctly labeled ones by using sub-centers
in face recognition tasks. Instead of summing sub-center logits in the final
output as previous works, they select the maximum sub-center logit and
disregard the others to have highly distinct sub-centers for correct and
incorrect class members.

The sub-center classification with noisy-label aspect is commonly applied in
speaker recognition from unlabeled data [20, 21]. However, the commonly used
sub-center modeling in image-retrieval tasks which applies summation of sub-
center logits, is fundamentally different from the noisy-label aspect of sub-
centers. Former type of sub-center modeling seems to be overlooked in speaker
recognition approaches. In this work, we adopt sub-center modeling from fine-
grained image retrieval tasks to speaker embeddings. We demonstrate its
effectiveness in better modeling speaker identity in the context of multi-
speaker, zero-shot speech synthesis.

##  III Sub-center Speaker Embeddings

In this work, we build upon one of the state-of-the-art speaker embedding
methods: Emphasized Channel Attention, Propagation and Aggregation in TDNN or
shortly ECAPA-TDNN [2]. ECAPA-TDNN is trained on a large amount of labeled
data with speaker classification objective and has been very successful in
speaker recognition. We introduce sub-class centers to ECAPA-TDNN by modifying
the output layer with the objective of making it more suitable for speech
synthesis.

###  III-A ECAPA-TDNN Speaker Embeddings

ECAPA-TDNN is a deep neural network architecture based on TDNN, which is
popular in many speech applications. As the main block, the model employs
squeeze-and-excitation Res2Net blocks together with TDNN layers. The model
also leverages channel-dependent attentive statistical pooling to distill and
focus on the important knowledge going through the deep layers. On top of
stacked main blocks, the model uses temporal pooling to gather information
from frame-level activations to a fixed-size vector, which is the speaker
embedding. In the training phase, after temporal pooling, the speaker
embedding is fed to a classifier head, and the model is trained with additive
angular margin softmax loss with speaker labels. In the network, the output
layer is $W\in\mathbb{R}^{L\times N}$ where $L$ is the number of hidden units
and $N$ is the number of output classes. Basically, the output layer consists
of $N$ class centers $w_{n}\in\mathbb{R}^{L}$.

The AAM-Softmax objective can be defined as

$l=-\log\frac{e^{s\cos(\theta_{x_{i}}+m)}}{e^{s\cos(\theta_{x_{i}}+m)}+\sum_{j=1,j\neq
x_{i}}^{N}e^{s\cos(\theta_{j})}}$ (1)

where $\theta_{x_{i}}=\arccos(w_{i}^{T}x_{i})$ is the angle between the
$i^{th}$ embedding $x_{i}$ and the ground truth class center $w_{i}$.
$\theta_{j}=\arccos(w_{j}^{T}x_{i})$ is the angle between the embedding
$x_{i}$ and other class centers, and $l$ corresponds to the negative log-
probability of $x_{i}$ being a member of the ground truth class $i$. $m$ is
the margin enforced between the correct class and other classes, and $s$ is
the scalar.

Figure 1: Proposed sub-center modeling (pink) on ECAPA-TDNN network

###  III-B Sub-center Modeling

We incorporate the sub-center modeling in [14] to AAM-Softmax objective of
ECAPA-TDNN. In order to have multiple sub-centers for each class, we modify
the output layer such that the weights become $W_{c}\in\mathbb{R}^{L\times
N\times C}$ where $C$ is the number of subcenters. For each class, we have a
set of sub-centers $\\{w_{n,1},w_{n,2},...,w_{n,c}\\}\in\mathbb{R}^{L}$. The
angle $\theta_{i}$ in standard AAM-Softmax becomes $G_{z_{i}}$ which is
calculated over sub-centers.

$l=-log\frac{e^{scos(G_{x_{i}}+m)}}{e^{scos(G_{x_{i}}+m)}+\sum_{j=1,j\neq
x_{i}}^{N}e^{scos(G_{j})}}$ (2)

We define $G(x_{i})$ as

$G_{x_{i}}=arccos(\sum_{c=1}^{C}\frac{\frac{1}{T}w_{i,c}^{T}x_{i}}{\sum_{c=1}^{C}\frac{1}{T}w_{i,c}^{T}x_{i}})$
(3)

where the angle is calculated from the summation of individual dot products
between sub-class centers $w_{i,c}$ and the embedding $x_{i}$, after passing
through the softmax function. with this approach, the embedding $x_{i}$ does
not have to be close to a single class center $w_{n}$. It can practically
select among the set of subcenters $\\{w_{n,1},w_{n,2},...,w_{n,c}\\}$ and be
close to one or multiple of them. $T$ is the temperature parameter. The
overall diagram of sub-center ECAPA-TDNN can be seen in Fig. 1.

Multiple sub-centers can capture the local variations in the utterances of a
speaker and introduce variance to embeddings since now embeddings from various
utterances do not have to collapse into a single center. This feature is
especially beneficial for speech synthesis tasks when modeling the output
speech distribution of the given speaker. Since all the sub-centers are still
subjected to the classification objective all the time, discriminative power
should also be preserved.

##  IV Voice Conversion with Speaker Embeddings

In our study, we concentrate on VC [22], which is a sub-task of speech
synthesis. VC focuses on altering the speaker identity of a given speech while
preserving its linguistic content. We applied our proposed speaker embeddings
to a state-of-the-art (SOTA) encoder-decoder-based voice conversion method
[10], where the speech is encoded to linguistic content, pitch, and speaker
representations by pre-trained encoders and decoded back to speech with a
Hifi-GAN vocoder [23] in the training phase. At inference, the method performs
voice conversion by extracting linguistic and pitch embeddings from the source
speech, while extracting speaker embeddings from a reference speech from the
target speaker. The output speech is decoded as the target speaker’s voice
uttering the source linguistic content.

###  IV-A Linguistic Representations

For the linguistic representation, the VC method uses discrete HuBERT units
[10, 24]. Discrete HuBERT units are obtained by applying k-means clustering to
the frame-level, continuous HuBERT features. Due to the automatic speech
recognition (ASR)-focused objective of masked unit prediction, linguistic
information is expected to be dominant in HuBERT features. HuBERT takes raw
waveform as input, and we use the 6th layer frame-level outputs from the model
for clustering. We use a k-means clustering model with K=100, trained on the
LibriSpeech-clean-100 [25] dataset.

###  IV-B Pitch Representations

The method also encodes normalized pitch contour into discrete pitch units.
The pitch sequence is extracted by Dio[26] algorithm with 20 ms window and 5
ms shift. The pitch contour is normalized using the pitch statistics of the
corresponding speaker.The normalized pitch contour is fed into a VQ-VAE based
encoder-decoder architecture [27]. The VQ-codebook indices of the encoded
frame-level f0 features $\mathbf{p}=[p_{1},p_{2},...,p_{T}]$ are used as the
discrete pitch units. The VQ-VAE model is trained on the VCTK dataset [28].

###  IV-C Speaker Representations

As the speaker representation, we employ our proposed sub-center ECAPA-TDNN
embeddings explained in Section 3. We extract the speaker embedding vector
$s\in\mathbb{R}^{192\times 1}$ from the short input utterance. The speaker
embedding vectors are normalized to unit length before being fed into the
decoder.

###  IV-D Vocoder

We use HifiGAN vocoder [23] to synthesize speech waveform from the linguistic,
pitch, and speaker representations. Only the dimensions of the original
HifiGAN implementation are modified to enable the use of our representations
as input. The discrete pitch sequence $\mathbf{p}=[p_{1},p_{2},...,p_{T_{p}}]$
is upsampled to match the length of the discrete HuBERT unit sequence
$\mathbf{h}=[h_{1},h_{2},...,h_{T_{h}}]$. The discrete sequences use
corresponding embedding lookup tables before the model and become
$\mathbf{h}_{emb}$ and $\mathbf{p}_{emb}$. Finally, the speaker embedding
vector $s$ is duplicated as many times as the length of $\mathbf{h}$, and the
concatenation of $z=[p_{emb},h_{emb},s_{dup}]$ is used as the input to the
vocoder.

##  V Experimental Setup

###  V-A Datasets

For the speaker embedding training, we used the VoxCeleb2 dataset [29]. For
voice conversion, we employed the VCTK corpus. VCTK contains a total of 110
speakers with 400 utterances from each speaker. For the VC experiments, 90
speakers are randomly chosen for training, while the remaining 20 speakers are
utilized for zero-shot VC experiments as unseen speakers.

###  V-B Training

The baseline standard ECAPA-TDNN is trained with the implementation and recipe
provided by SpeechBrain [30]. We implemented the sub-center ECAPA-TDNN on top
of that implementation and trained with the default parameters. In the
training, Adam optimizer is used with a base learning rate of 1e-4 and a
cyclic schedule. The batch size is 32, and online augmentation (noise
addition, reverberation) is applied to batch utterances as in [2]. The margin
$m$ and scale $s$ parameters for the AAM-Softmax loss are $0.4$ and $30$,
respectively. For the VC training, we used the implementation in speech-
resynthesis [10]. We modified the original implementation to incorporate
ECAPA-TDNN embeddings as the speaker representations. All the encoders for
linguistic, f0, and speaker representations are trained independently and
beforehand. The HifiGAN vocoder is trained while those encoders explained in
Section 3 are used as pre-trained and remain frozen. We used default
parameters as in [10].

###  V-C Evaluation

####  V-C1 Speaker Verification & Intra-class Variance

Firstly, we measured the intra-class variance of standard ECAPA-TDNN and our
proposed embeddings. Since we are comparing different embedding spaces, we
calculated the ratio of intra-class variance to inter-class as the nominal
intra-class variance measure. We calculate intra-class variance as

$\sigma_{intra-
class}^{2}=\frac{\sum^{N}(f(x_{s,i},\tilde{x_{s}})-\mu_{intra})^{2}}{N}$ (4)

where $x_{s,i}$ is the $i^{t}h$ speaker embedding from speaker $s$,
$\tilde{x_{s}}$ is the mean of all embeddings from speaker $s$, and $f$ is the
cosine similarity function. $\mu_{intra}$ is the mean off all intra-class
cosine distances and $N$ is the total number of examples. We calculate the
inter-class variance as

$\sigma_{\text{inter-
class}}^{2}=\frac{\sum^{N\times(S-1)}(f(x_{s,i},\tilde{x_{s^{\prime}}})-\mu_{\text{inter}})^{2}}{N\times(S-1)}$
(5)

TABLE I: Speaker Recognition and Intra-class variance results

Embedding EER(%) var ECAPA-TDNN [2] 1.71 0,42 Sub-center ECAPA-TDNN C = 10
1.50 0,45 Sub-center ECAPA-TDNN C = 10, T = 0.1 1.47 0,36 Sub-center ECAPA-
TDNN, C = 20 1.55 0,47

TABLE II: Objective evaluations for VC

Method WER $\downarrow$ CER $\downarrow$ SECS $\uparrow$ VC with ECAPA-TDNN
[2] 14.84 6.82 64,04 VC with Sub ECAPA-TDNN, C = 10 14.65 6.72 64,14 VC with
Sub ECAPA-TDNN, C = 10, T = 0.1 14,32 6,67 65,86 VC with Sub ECAPA-TDNN, C =
20 13.93 6.41 64.59

where we measure the distance between the $i$-th speaker embedding from
speaker $s$ and every other speaker’s mean embedding $s^{\prime}$ different
from $s$. As the final inter-class variation measure, we report the ratio
$\sigma_{\text{intra-class}}^{2}/\sigma_{\text{inter-class}}^{2}$.

For the speaker verification evaluation, we created a total of 20M trials from
110 speakers of VCTK dataset and measured the equal-error-rate (EER) on cosine
similarity scores of positive and negative pairs.

####  V-C2 VC Evaluation

In VC experiments, we performed both objective and subjective evaluations. For
the objective evaluations, we measured word error rate (WER) and character
error rate (CER) when the synthesized speech is fed into SOTA ASR
model44footnotemark: 4 [31]. We also measured speaker embedding cosine
similarity (SECS) between embeddings from synthesized speech and ground truth
utterances using the d-vector111https://github.com/resemble-ai/Resemblyzer[3].
In total, we obtained around 20k converted utterances between test speaker
pairs. For the subjective evaluation, we conducted mean-opinion-score (MOS)
[32] for naturalness, similarity mean-opinion-score [4] (SMOS) for speaker
similarity and ABX preference tests [32] for more consistent prosody (e.g
intonation, stress, rhythm) evaluation. 120 converted speech samples were
presented to 12 participants.

##  VI Results

22footnotetext: https://huggingface.co/facebook/wav2vec2-large-960h-lv60-self

###  VI-A Speaker Verification & Intra-class Variance

As the number of sub-centers for each class, we experimented with $C=10,20$
following the work in [14]. We also tried with no temperature ($T=1$) and a
small temperature ($T=0.1$) in the sub-center logit calculation. We reported
speaker verification EER and intra-class/inter-class distance variance ratio
as $var$ in Table II. Results show that sub-center modeling with no
temperature has higher intra-class variance than the standard ECAPA-TDNN.
Moreover, although the sub-center models have higher intra-class variance,
they do not lose any discriminative power as they are also better in terms of
verification EER, which indicates a better modeling of complex intra-class
distributions.

Most surprisingly, sub-center modeling with a low temperature ($T=0.1$) shows
a lower intra-class variance than the baseline. We believe a low-temperature
value greatly affects the sub-center selection by making the selections
extremely confident, which results in utilizing few sub-centers during
training which is also addressed by previous sub-center works [15, 18]. We
show that the temperature parameter is one of the tools for adjusting the
utilization of sub-centers. As expected, sub-center modeling with the lowest
intra-class variance has the best verification performance.

TABLE III: Subjective Evaluation Results for VC in 95% confidence interval

Method MOS SMOS Ground Truth 4.65 $\pm$ 0.09 - VC with ECAPA-TDNN [2] 2.94
$\pm$ 0.12 2.65 $\pm$ 0.13 VC with Sub-center ECAPA-TDNN, C = 10, T = 0.1 2.89
$\pm$ 0.13 2.76 $\pm$ 0.13 VC with Sub-center ECAPA-TDNN, C = 20 3.18 $\pm$
0.12 2.88 $\pm$ 0.13

Figure 2: ABX prosody preference test results between the VC with speaker
embeddings having highest and lowest intra-class variance

###  VI-B Voice Conversion

The objective evaluation results can be seen in Table II. The embedding with
the highest intra-class variance (sub-center ECAPA-TDNN, C=20) shows the
lowest WER/CER indicating a better intelligibility and higher quality of
synthesized speech. It also improves SECS compared to baseline ECAPA-TDNN
embedding. The embedding with the lowest intra-class variance (sub-center
ECAPA-TDNN, C=10, T=0.1) shows best SECS. Comparing embeddings with the
highest and lowest intra-class variance, the one with higher variance has
better WER/CER, while the one with lower variance has better SECS, though they
have comparable SV performance. We note that both of them achieve better
results than the baseline in terms of WER/CER and SECS.

Subjective evaluation results in Table III and Fig. 2 show that embeddings
with the highest intra-class variance are the best for naturalness, proving
the effectiveness of introducing variation to speaker embedding. We applied a
one-tailed, paired t-test to the MOS results, and with $p<0.05$, the
statistical significance is affirmed. The proposed embedding with higher
intra-class variance also performs best in similarity MOS results,
demonstrating its effectiveness in terms of speaker similarity. In the ABX
preference test for prosody, embeddings with the highest variance outperform
those with the lowest variance. The proposed sub-center modeling improves upon
the baseline in almost all results. While lower intra-class variance offers
better speaker discrimination than the baseline, higher variance embeddings
are significantly better at capturing voice variations and synthesizing the
most natural speech with the highest speaker similarity.

##  VII Conclusion

Originally designed for speaker recognition, speaker embeddings are widely
used in synthesis tasks due to their advantages. Our work highlights the
discrepancies between speaker recognition and speech synthesis, introducing
speaker embeddings that are more suitable for synthesis. We propose
incorporating variation information into speaker embeddings which enhances
speaker discriminability at the same time. Applying our proposed embeddings to
the VC task shows that those with higher intra-class variance perform best in
naturalness and speaker similarity, as they better model the real speech
distribution. Additionally, prosody-focused evaluations demonstrate that the
proposed embeddings excel in modeling speech variations.

## References

  * [1] D. Snyder, D. Garcia-Romero, G. Sell, D. Povey, and S. Khudanpur, “X-vectors: Robust dnn embeddings for speaker recognition,” in _2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2018, pp. 5329–5333.
  * [2] B. Desplanques, J. Thienpondt, and K. Demuynck, “ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in TDNN Based Speaker Verification,” in _Proc. Interspeech 2020_ , 2020, pp. 3830–3834.
  * [3] L. Wan, Q. Wang, A. Papir, and I. Lopez-Moreno, “Generalized end-to-end loss for speaker verification,” _2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 4879–4883, 2017. [Online]. Available: https://api.semanticscholar.org/CorpusID:22987563
  * [4] E. Casanova, J. Weber, C. D. Shulby, A. C. Junior, E. Gölge, and M. A. Ponti, “YourTTS: Towards zero-shot multi-speaker TTS and zero-shot voice conversion for everyone,” in _Proceedings of the 39th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, Eds., vol. 162. PMLR, 17–23 Jul 2022, pp. 2709–2720. [Online]. Available: https://proceedings.mlr.press/v162/casanova22a.html
  * [5] E. Cooper, C.-I. Lai, Y. Yasuda, F. Fang, X. E. Wang, N. Chen, and J. Yamagishi, “Zero-shot multi-speaker text-to-speech with state-of-the-art neural speaker embeddings,” _ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 6184–6188, 2019. [Online]. Available: https://api.semanticscholar.org/CorpusID:204852286
  * [6] K. Qian, Y. Zhang, S. Chang, X. Yang, and M. Hasegawa-Johnson, “AutoVC: Zero-shot voice style transfer with only autoencoder loss,” in _Proceedings of the 36th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, K. Chaudhuri and R. Salakhutdinov, Eds., vol. 97. PMLR, 09–15 Jun 2019, pp. 5210–5219. [Online]. Available: https://proceedings.mlr.press/v97/qian19c.html
  * [7] H. Siuzdak, P. Dura, P. van Rijn, and N. Jacoby, “WavThruVec: Latent speech representation as intermediate features for neural speech synthesis,” in _Proc. Interspeech 2022_ , 2022, pp. 833–837.
  * [8] J. Zhang, S. Jayasuriya, and V. Berisha, “Learning repeatable speech embeddings using an intra-class correlation regularizer,” in _Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023_ , A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, Eds., 2023. [Online]. Available: http://papers.nips.cc/paper\\_files/paper/2023/hash/f0aa7e9e67515fa0c607c2959ccda6a0-Abstract-Conference.html
  * [9] J. Li, W. Tu, and L. Xiao, “Freevc: Towards high-quality text-free one-shot voice conversion,” in _ICASSP 2023 - 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2023, pp. 1–5.
  * [10] A. Polyak, Y. Adi, J. Copet, E. Kharitonov, K. Lakhotia, W.-N. Hsu, A. Mohamed, and E. Dupoux, “Speech Resynthesis from Discrete Disentangled Self-Supervised Representations,” in _Proc. Interspeech 2021_ , 2021, pp. 3615–3619.
  * [11] V. Popov, I. Vovk, V. Gogoryan, T. Sadekova, M. S. Kudinov, and J. Wei, “Diffusion-based voice conversion with fast maximum likelihood sampling scheme,” in _International Conference on Learning Representations_ , 2022. [Online]. Available: https://openreview.net/forum?id=8c50f-DoWAu
  * [12] N. Le and J.-M. Odobez, “Robust and Discriminative Speaker Embedding via Intra-Class Distance Variance Regularization,” in _Proc. Interspeech 2018_ , 2018, pp. 2257–2261.
  * [13] J. S. Chung, J. Huh, S. Mun, M. Lee, H.-S. Heo, S. Choe, C. Ham, S. Jung, B.-J. Lee, and I. Han, “In Defence of Metric Learning for Speaker Recognition,” in _Proc. Interspeech 2020_ , 2020, pp. 2977–2981.
  * [14] Q. Qian, L. Shang, B. Sun, J. Hu, T. Tacoma, H. Li, and R. Jin, “Softtriple loss: Deep metric learning without triplet sampling,” in _2019 IEEE/CVF International Conference on Computer Vision (ICCV)_. Los Alamitos, CA, USA: IEEE Computer Society, nov 2019, pp. 6449–6457. [Online]. Available: https://doi.ieeecomputersociety.org/10.1109/ICCV.2019.00655
  * [15] Z. Zhang and X. Gong, “The fixed sub-center: A better way to capture data complexity,” _ArXiv_ , vol. abs/2203.12928, 2022.
  * [16] E. Levi, T. Xiao, X. Wang, and T. Darrell, “Rethinking preventing class-collapsing in metric learning with margin-based losses,” in _2021 IEEE/CVF International Conference on Computer Vision (ICCV)_. Los Alamitos, CA, USA: IEEE Computer Society, oct 2021, pp. 10 296–10 305. [Online]. Available: https://doi.ieeecomputersociety.org/10.1109/ICCV48922.2021.01015
  * [17] P. Liu, Z. Liu, Y. Lang, S. Liu, Q. Zhou, and Q. Li, “Deep metric learning assisted by intra-variance in a semi-supervised view of learning,” _Engineering Applications of Artificial Intelligence_ , vol. 131, p. 107885, 2024. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S0952197624000435
  * [18] R. Müller, S. Kornblith, and G. E. Hinton, “Subclass distillation,” _ArXiv_ , vol. abs/2002.03936, 2020.
  * [19] J. Deng, J. Guo, T. Liu, M. Gong, and S. Zafeiriou, “Sub-center arcface: Boosting face recognition by large-scale noisy web faces,” in _Computer Vision – ECCV 2020_ , A. Vedaldi, H. Bischof, T. Brox, and J.-M. Frahm, Eds. Cham: Springer International Publishing, 2020, pp. 741–757.
  * [20] J. Thienpondt, B. Desplanques, and K. Demuynck, “The idlab voxceleb speaker recognition challenge 2020 system description,” _ArXiv_ , vol. abs/2010.12468, 2020. [Online]. Available: https://api.semanticscholar.org/CorpusID:260516427
  * [21] Y. Zheng, J. Peng, Y. Chen, Y. Zhang, J. Wang, M. Liu, and M. Xu, “The speakin speaker verification system for far-field speaker verification challenge 2022,” _ArXiv_ , vol. abs/2209.11625, 2022. [Online]. Available: https://api.semanticscholar.org/CorpusID:252519303
  * [22] B. Sisman, J. Yamagishi, S. King, and H. Li, “An overview of voice conversion and its challenges: From statistical modeling to deep learning,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , vol. 29, pp. 132–157, 2021.
  * [23] J. Kong, J. Kim, and J. Bae, “Hifi-gan: Generative adversarial networks for efficient and high fidelity speech synthesis,” in _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, Eds., vol. 33. Curran Associates, Inc., 2020, pp. 17 022–17 033. [Online]. Available: https://proceedings.neurips.cc/paper_files/paper/2020/file/c5d736809766d46260d816d8dbc9eb44-Paper.pdf
  * [24] W.-N. Hsu, B. Bolte, Y.-H. H. Tsai, K. Lakhotia, R. Salakhutdinov, and A. Mohamed, “Hubert: Self-supervised speech representation learning by masked prediction of hidden units,” _IEEE/ACM Trans. Audio, Speech and Lang. Proc._ , vol. 29, p. 3451–3460, oct 2021. [Online]. Available: https://doi.org/10.1109/TASLP.2021.3122291
  * [25] V. Panayotov, G. Chen, D. Povey, and S. Khudanpur, “Librispeech: An asr corpus based on public domain audio books,” in _2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , 2015, pp. 5206–5210.
  * [26] M. MORISE, F. YOKOMORI, and K. Ozawa, “World: A vocoder-based high-quality speech synthesis system for real-time applications,” _IEICE Transactions on Information and Systems_ , vol. E99.D, pp. 1877–1884, 07 2016.
  * [27] P. Dhariwal, H. Jun, C. Payne, J. W. Kim, A. Radford, and I. Sutskever, “Jukebox: A generative model for music,” 2020.
  * [28] J. Yamagishi, C. Veaux, and K. MacDonald, “CSTR VCTK Corpus: English multi-speaker corpus for CSTR voice cloning toolkit (version 0.92),” 2019.
  * [29] J. S. Chung, A. Nagrani, and A. Zisserman, “Voxceleb2: Deep speaker recognition,” in _Interspeech_ , 2018. [Online]. Available: https://api.semanticscholar.org/CorpusID:49211906
  * [30] M. Ravanelli, T. Parcollet, P. Plantinga, A. Rouhe, S. Cornell, L. Lugosch, C. Subakan, N. Dawalatabad, A. Heba, J. Zhong, J.-C. Chou, S.-L. Yeh, S.-W. Fu, C.-F. Liao, E. Rastorgueva, F. Grondin, W. Aris, H. Na, Y. Gao, R. D. Mori, and Y. Bengio, “SpeechBrain: A general-purpose speech toolkit,” 2021, arXiv:2106.04624.
  * [31] A. Baevski, H. Zhou, A. Mohamed, and M. Auli, “wav2vec 2.0: a framework for self-supervised learning of speech representations,” in _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , ser. NIPS ’20. Red Hook, NY, USA: Curran Associates Inc., 2020.
  * [32] B. Sisman, J. Yamagishi, S. King, and H. Li, “An overview of voice conversion and its challenges: From statistical modeling to deep learning,” _IEEE/ACM Trans. Audio, Speech and Lang. Proc._ , vol. 29, p. 132–157, nov 2020. [Online]. Available: https://doi.org/10.1109/TASLP.2020.3038524


# Randomized Value Functions via
Posterior State-Abstraction Sampling

Dilip Arumugam<EMAIL_ADDRESS>Stanford University  Benjamin Van Roy
<EMAIL_ADDRESS>Stanford University  DeepMind

###### Abstract

State abstraction has been an essential tool for dramatically improving the
sample efficiency of reinforcement-learning algorithms. Indeed, by exposing
and accentuating various types of latent structure within the environment,
different classes of state abstraction have enabled improved theoretical
guarantees and empirical performance. When dealing with state abstractions
that capture structure in the value function, however, a standard assumption
is that the true abstraction has been supplied or unrealistically computed a
priori, leaving open the question of how to efficiently uncover such latent
structure while jointly seeking out optimal behavior. Taking inspiration from
the bandit literature, we propose that an agent seeking out latent task
structure must explicitly represent and maintain its uncertainty over that
structure as part of its overall uncertainty about the environment. We
introduce a practical algorithm for doing this using two posterior
distributions over state abstractions and abstract-state values. In
empirically validating our approach, we find that substantial performance
gains lie in the multi-task setting where tasks share a common, low-
dimensional representation.

††footnotetext: Preprint. Work in progress.

##  1 Introduction

The complexity of the state space within an environment has profound impact on
a decision-making agent’s capacity for sample-efficient reinforcement learning
(Kearns and Singh, 2002; Kakade, 2003; Strehl et al., 2009; Auer et al.,
2009). Oftentimes, however, this complexity (as measured by the cardinality of
a finite set of states or the dimensionality of real-valued
observations/features) is an exaggeration of the amount of information
actually needed to make optimal (or even near-optimal) decisions. And yet,
many reinforcement-learning algorithms make no concerted effort to fully
exploit this structure so as to accelerate learning of the optimal policy or
value function (Watkins and Dayan, 1992; Williams, 1992; Sutton, 1988; Sutton
et al., 2000; Mnih et al., 2015).

In contrast, there exist several works which closely examine state
abstractions from a theoretical or empirical perspective, highlighting the
advantages of acknowledging and utilizing latent problem structure (Bertsekas
et al., 1988; Dean and Givan, 1997; Ferns et al., 2004; Jong and Stone, 2005;
Li et al., 2006; Van Roy, 2006; Ferns et al., 2012; Jiang et al., 2015; Abel
et al., 2016, 2019; Dong et al., 2019; Du et al., 2019; Misra et al., 2019;
Russo, 2020). While important for gaining clarity and an appreciation for what
state abstraction brings to the reinforcement-learning problem (Lake et al.,
2017; Konidaris, 2019), some of these works can be unrealistic in their
acquisition of the very object being studied. Specifically, those approaches
focusing on abstractions that capture structure within the optimal value
function often assume that an ideal state abstraction has been provided by an
oracle (Li et al., 2006; Abel et al., 2016; Dong et al., 2019); while this may
be a suitable provision for the purpose of theoretical analysis, this
engenders a two-step procedure in practice whereby the optimal value function
is computed exactly and then used to compute the corresponding state
abstraction (Abel et al., 2016, 2018). These approaches beg a natural
question: can an agent learn a value-based state abstraction while
simultaneously learning how to act optimally within an environment?
Incorporating this additional abstraction criterion, beyond the standard RL
objective of learning an optimal policy, amounts to inserting a form of
inductive bias within the learning algorithm. Thus, a natural follow-up
question arises: when is it advantageous to learn a state abstraction
alongside learning optimal behavior?

In this work, we offer an affirmative answer to the first question and provide
an empirically-supported hypothesis to the second. In addressing the former,
we build upon an existing line of work that grounds (both theoretically and
empirically) the efficient learning of optimal behavior through an agent’s
representation of epistemic uncertainty over its environment (Russo and Van
Roy, 2016; Osband et al., 2016a; Russo and Van Roy, 2018; O’Donoghue et al.,
2018; Osband et al., 2019). Informally, it is the agent’s resolution of this
uncertainty that naturally fosters deep exploration and provably-efficient
learning. Adopting this perspective, we assert that an agent must maintain an
explicit belief over the state abstraction that underlies the environment.
This belief can then be coupled with a conditional posterior over abstract-
state values to render the agent’s full uncertainty over the optimal value
function. As for when it is advantageous to incorporate this factored
posterior, we empirically demonstrate considerable performance gains when an
agent engages with multiple tasks in the same environment, all supported by a
single, underlying state abstraction.

The paper proceeds as follows: in Section 2, we formulate the problem and
clarify our setting where the optimal value function admits a natural state
abstraction, which we formalize as exhibiting low-rank structure. We then
discuss related work and background in Sections 3 and 4, respectively, before
introducing our approach, Posterior State-abstraction Sampling (PS2), in
Section 5. We conclude with illustrative experiments in contextual-bandit
problems, exploring both the single-task and multi-task learning settings
(Sections 6 and 7).

##  2 Problem Formulation

###  2.1 Reinforcement Learning

We formulate an individual task as a finite-horizon, episodic Markov Decision
Process (MDP) (Bellman, 1957; Puterman, 1994) defined by
$\mathcal{M}=\langle\mathcal{S},\mathcal{A},\mathcal{R},\mathcal{T},\gamma,H\rangle$
where $\mathcal{S}$ denotes the state space, $\mathcal{A}$ is the finite
action set, $\mathcal{R}_{h}:\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}$ is
a (deterministic) reward function for timestep $h$,
$\mathcal{T}:\mathcal{S}\times\mathcal{A}\mapsto\Delta(\mathcal{S})$ is the
transition function producing a distribution over next states given the
current state-action pair, $\gamma\in[0,1)$ is the discount factor, and
$H\in\mathbb{N}$ is the horizon or episode duration. We use $|\mathcal{A}|=A$
and, when $\mathcal{S}$ is finite, $|\mathcal{S}|=S$ to denote the respective
sizes of the action space and state space.

Learning proceeds in $T$ stages or episodes where, at each timestep of the
current episode $h\in[H]$, the agent observes the current state $s_{h}$ and
selects an action $a_{h}$ according to its current policy
$\pi_{h}:\mathcal{S}\mapsto\Delta(\mathcal{A})$. We assume that $\mathcal{M}$
has a fixed initial state distribution $\rho\in\Delta(\mathcal{S})$ such that
$s_{1}\sim\rho(\cdot)$ is observed at the start of each episode. The objective
for the agent is to synthesize a policy so as to maximize expected return
$\mathbb{E}[\sum\limits_{h=1}^{H}\gamma^{h-1}\mathcal{R}_{h}(s_{h},a_{h})]$.
The value function of a policy $\pi$ denotes the expected future discounted
return by following the policy from a given state $s$:

$\displaystyle
V_{h}^{\pi}(s_{h})=\mathbb{E}[\sum\limits_{h^{\prime}=0}^{H-h}\gamma^{h^{\prime}}\mathcal{R}_{h}(s_{h^{\prime}},a_{h^{\prime}})|s_{h^{\prime}}=s_{h}]$

where the expectation is taken with respect to the stochasticity in the
environment transition dynamics and policy. Similarly, we use the action-value
function $Q^{\pi}(s,a)$ to define the expected future discounted return from
being in a state $s$, taking action $a$, and following policy $\pi$
thereafter:

$\displaystyle
Q_{h}^{\pi}(s,a)=\mathcal{R}_{h}(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim\mathcal{T}(\cdot|s,a)}[V_{h+1}^{\pi}(s^{\prime})]$

Denoting the policy class containing all stationary, stochastic policies as
$\Pi=\\{\pi|\pi:\mathcal{S}\mapsto\Delta(\mathcal{A})\\}$, we may define the
optimal policy $\pi^{\star}=\sup\limits_{\pi\in\Pi^{H}}V_{1}^{\pi}(s)$, where
$\Pi^{H}$ denotes the class of non-stationary policies, whose value functions
are given by the Bellman optimality equations:

$\displaystyle V_{h}^{\star}(s)$
$\displaystyle=\max\limits_{a\in\mathcal{A}}Q_{h}^{\star}(s,a)$ $\displaystyle
Q_{h}^{\star}(s,a)$
$\displaystyle=\mathcal{R}_{h}(s,a)+\gamma\mathbb{E}_{s^{\prime}\sim\mathcal{T}(\cdot|s,a)}[\max\limits_{a^{\prime}\in\mathcal{A}}Q_{h+1}^{\star}(s^{\prime},a^{\prime})]$

where $V_{H+1}^{\star}(s)=0$. Concretely, our goal is to leverage experience
sampled in each episode to learn $Q^{\star}$, yielding the optimal policy
$\pi_{h}^{\star}(s)=\operatorname*{arg\,max}\limits_{a\in\mathcal{A}}Q_{h}^{\star}(s,a)$
(Sutton and Barto, 1998; Watkins and Dayan, 1992).

###  2.2 State Abstractions

A well-studied tool for accelerating RL algorithms is the use of state
aggregation or state abstraction to reduce the size of the MDP state space
(Bertsekas et al., 1988; Li et al., 2006; Van Roy, 2006). Indeed, given the
dependence on $|\mathcal{S}|$ that appears in numerous sample complexity
results for RL (Kakade, 2003; Strehl et al., 2009), it seems only natural that
a reduction in the overall number of states under consideration can
dramatically alleviate the burdens of learning an optimal policy.

As outlined in Li et al. (2006), several classes of state abstractions exist,
each elucidating structure contained in different components of the MDP. A
choice of state abstraction from one of these classes characterizes a
particular function $\phi:\mathcal{S}\mapsto\mathcal{S}_{\phi}$ mapping
original or ground states of the MDP into an aggregate or abstract state space
$\mathcal{S}_{\phi}$. Naturally, the abstract state space is taken to be
smaller, in some sense, than the original (for instance,
$|\mathcal{S}_{\phi}|<|\mathcal{S}|$) such that $\phi$ defines a (lossy)
compression of the original state space (Abel et al., 2019). While state
abstraction constitutes a general mechanism for specifying latent task
structure, several works often make generous assumptions concerning the
provision or acquisition of $\phi$ in order to highlight the benefits of
leveraging such structure in RL. In contrast, this work is concerned with
weakening those assumptions specifically by learning $\phi$ concurrently with
$Q^{\star}$.

Various prior works have focused on the capacity of state abstraction to
enable provably-efficient RL algorithms that scale to tasks with high-
dimensional observations (Du et al., 2019; Misra et al., 2019; Agarwal et al.,
2020). These approaches build state abstractions based on latent structure in
the transition dynamics; by operating in the reward-free setting (Hazan et
al., 2019; Jin et al., 2020), which lacks a specific reward function, these
algorithms employ various mechanisms to facilitate thorough exploration of the
environment, yielding a strong approximation of the transition model. This
style of approach seems like a natural “path of least resistance” in so far as
each step within the environment yields a ground-truth signal that can be
aimed at further distilling the true, underlying abstraction. In contrast, a
state abstraction based on the similarity of $Q^{\star}$-values (Li et al.,
2006; Abel et al., 2016, 2019; Dong et al., 2019) inherently draws upon
knowledge of $Q^{\star}$ which, if readily accessible to the agent, would
imply knowledge of the optimal policy itself. Nevertheless, in this paper, we
explicitly direct our focus to these $Q^{\star}$-based state abstractions and
avert the apparent “chicken-or-egg” problem through the agent’s own epistemic
uncertainty about the environment.

###  2.3 Low-Rank Value Functions

We formalize the type of $Q^{\star}$-based state abstraction studied in this
work under the following definition of a low-rank value function:  Definition
1 (Low-Rank $Q^{\star}$-function):  The optimal action-value function
$Q^{\star}$ of an MDP $\mathcal{M}$ is characterized as low-rank if there
exists two functions, $\phi_{h}^{\star}:\mathcal{S}\mapsto\mathbb{R}^{M}$ and
$\psi_{h}^{\star}:\mathcal{A}\mapsto\mathbb{R}^{M}$, such that $\forall
s\in\mathcal{S},a\in\mathcal{A},h\in[H]$: $\displaystyle Q_{h}^{\star}(s,a)$
$\displaystyle=\langle\phi_{h}(s),\psi_{h}(a)\rangle$ where $\langle
u,v\rangle$ denotes the inner product between vectors $u,v$. We interpret the
latent dimension $M\in\mathbb{N}$ as the number of abstract states.

######  Assumption 1.

Throughout this work, we will assume that $\mathcal{M}$ admits a low-rank
$Q^{\star}$-function with rank $M$.

To clarify this structural assumption, consider the case where $\mathcal{M}$
has a finite state-action space. Dropping the timestep subscript for clarity,
note that $Q^{\star}$ may then be compactly represented as a matrix
$\mathbf{Q}^{\star}\in\mathbb{R}^{S\times A}$ where
$\mathbf{Q}_{ij}^{\star}=Q^{\star}(s_{i},a_{j})$, for some arbitrary indexing
of states and actions. Defining the rank of a matrix as  Definition 2 (Matrix
Rank):  For an arbitrary matrix $Y\in\mathbb{R}^{n\times m}$, the rank of $Y$
is defined as $\displaystyle\text{rank}(Y)=\min\\{k|k\in\mathbb{N}_{+},\exists
U\in\mathbb{R}^{n\times k},V\in\mathbb{R}^{m\times k}\text{ such that
}Y=UV^{T}\\}$ Assumption 1 then follows as
$\text{rank}(\mathbf{Q}^{\star})=M$. In words, we assume that the matrix of
optimal action-values admits a low-rank factorization
$\mathbf{Q}^{\star}=\mathbf{\Phi}\mathbf{\Psi}^{T}$, with
$\mathbf{\Phi}\in\mathbb{R}^{S\times M}$, $\mathbf{\Psi}\in\mathbb{R}^{A\times
M}$. Under this view, we can interpret learning a low-rank decomposition of
$\mathbf{Q}^{\star}$ as constructing predictive features in a $M$-dimensional
subspace that enable accurate estimation of $Q^{\star}$ for each state-action
pair (Tsitsiklis and Van Roy, 1996).

We conclude this section by briefly making explicit some connections between
low-rank value functions and the $Q^{\star}$-similarity state abstractions
employed in prior work. Specifically, notice that the exact
$Q^{\star}$-irrelevance abstraction of Li et al. (2006) corresponds to an
abstraction function $\phi^{\star}:\mathcal{S}\mapsto\\{0,1\\}^{M}$ where
$\forall s\in\mathcal{S},\sum\limits_{i=1}^{M}\phi(s)_{i}=1$; in words,
$\phi^{\star}$ maps each state to a one-hot vector thereby guaranteeing that
$\phi(s_{1})=\phi(s_{2})\iff Q^{\star}(s_{1},a)=Q^{\star}(s_{2},a),\forall
s_{1},s_{2}\in\mathcal{S},a\in\mathcal{A}$. An approximate version of this
abstraction, as studied in (Abel et al., 2016; Dong et al., 2019), can be
achieved by maintaining the same conditions on $\phi^{\star}$ and weakening
the definition of a low-rank value function to
$|Q^{\star}(s,a)-\langle\phi^{\star}(s)\psi^{\star}(a)\rangle|\leq\frac{\epsilon}{2}$,
for some constant $\epsilon$. Finally, the soft state aggregations of Singh et
al. (1995) impose the alternative condition
$\phi^{\star}:\mathcal{S}\mapsto\Delta^{M-1}$, mapping individual states to
the $(M-1)$-dimensional probability simplex.

##  3 Related Work

As previously discussed, this paper falls in with a long, rich line of work on
state abstraction in reinforcement learning (Whitt, 1978; Bertsekas et al.,
1988; Dean and Givan, 1997; Ferns et al., 2004; Jong and Stone, 2005; Li et
al., 2006; Van Roy, 2006; Ferns et al., 2012; Jiang et al., 2015; Abel et al.,
2016, 2019; Dong et al., 2019; Du et al., 2019; Misra et al., 2019). Notably,
this work is concerned with how an agent may incrementally learn a state
abstraction that capitalizes on latent structure in the optimal value function
(Bertsekas et al., 1988; Van Roy, 2006; Li et al., 2006; Abel et al., 2016,
2018, 2019; Dong et al., 2019). Bertsekas et al. (1988) focus on adaptively
synthesizing state aggregations based on Bellman-error residuals. Van Roy
(2006) examines and provides performance guarantees on approximate value
iteration under the provision of a particular state aggregation, leaving open
the question of how to dynamically abstract states and maintain performance
guarantees. Li et al. (2006) offer a unified perspective on a broad array of
state abstraction types, showcasing how state abstraction based on
$Q^{\star}$-similarity preserves the optimal policy. Later, Abel et al. (2016)
generalize this to the case of approximate state abstraction, highlighting the
approximation parameter as a knob for weighing state-space compression against
value loss. Abel et al. (2018) examine the lifelong learning setting where
$Q^{\star}$ is computed exactly for some number of MDPs before a
$Q^{\star}$-similarity state abstraction is then applied for the remainder of
the task distribution. Abel et al. (2019) formalize this intuition using tools
from rate-distortion theory (Shannon, 1959), but restrict focus to the
apprenticeship learning setting. Common to all of these works is the lack of a
practical, scalable algorithm for jointly learning the state abstraction and
corresponding abstract-state values, without knowledge of $Q^{\star}$; our
work rectifies this and offers one such approach.

Our work is also intimately related to the problem of low-rank matrix
completion or factorization. While the general problem of low-rank matrix
completion is underspecified and known to be NP-hard (Chistov and Grigor’Ev,
1984; Hardt et al., 2014), a large body of prior work identifies sufficient
conditions for designing provably-efficient factorization algorithms (Candès
and Recht, 2009; Candès and Tao, 2010; Candès and Plan, 2010; Keshavan et al.,
2010a, b; Recht, 2011). The algorithm presented in this work aligns with
gradient-descent based approaches for iteratively optimizing the latent
factors $\mathbf{\Phi},\mathbf{\Psi}$ which, despite their weaker sample
complexity guarantees (Jain et al., 2013; Gunasekar et al., 2013; Hardt, 2014;
Chen and Wainwright, 2015), are simple, scalable, and widely deployed in
practice (Recht and Ré, 2013; Koren et al., 2009). Unlike the standard
formulation of the low-rank matrix completion problem, our focus on the
sequential decision-making setting more closely aligns with adaptive-sampling
approaches to matrix completion (Krishnamurthy and Singh, 2013, 2014) which
are known to enjoy better sample complexity guarantees.

Several papers adopt a matrix factorization perspective for state abstraction
in reinforcement learning (Barreto et al., 2016; Behzadian and Petrik, 2018;
Duan et al., 2019; Zhang and Wang, 2019; Yang and Wang, 2019b, a; Agarwal et
al., 2020); crucially, however, these works opt for computing a factorization
of the transition function, rather than the value function. Again, we suspect
that this preference stems from the immediate inaccessibility of $Q^{\star}$,
a fact that we show need not be an obstacle when adopting a Bayesian view of
efficient exploration. Moreover, leveraging such a Bayesian approach to
exploration potentially avoids the known pitfalls (Osband and Van Roy, 2017)
of alternative methods that employ optimism in the face of uncertainty (Yang
and Wang, 2019b). Most related to this work is the approach of Shah et al.
(2020) who do in fact aim to learn the singular value decomposition of
$Q^{\star}$ and take advantage of low-rank structure; important differences
from this work include a focus on MDPs with continuous state-action spaces
along with assumptions on the Lipschitz continuity of $Q^{\star}$, access to a
generative model for sampling transitions, and access to anchor states (states
that are representative of each latent abstract state (Donoho and Stodden,
2004)) for heuristically guiding exploration/data collection. In contrast,
this work is concerned with discrete-action MDPs (though extensions to the
continuous-control setting are a natural future direction) and makes no
assumptions on $Q^{\star}$ beyond being low rank (Assumption 1).

For understanding when it is prudent for an agent to pursue latent task
structure in the form of a state abstraction, we find the multi-task setting
to be a natural candidate. In the context of the low-rank matrix factorization
outlined in the previous section, this amounts to asserting that the optimal
value functions of all tasks share a common, latent factor $\mathbf{\Phi}$
while each individual task $t$ also yields a specific matrix of abstract-state
values, $\mathbf{\Psi}_{t}$. Consequently, an agent interacting to solve all
tasks in parallel can greatly benefit from synthesizing shared task structure.
Various prior works already assess such benefits of multiple tasks for
supervised learning (Caruana, 1997; Baxter, 1997, 2000). While work that
formally develops this connection for reinforcement learning is still nascent
(D’Eramo et al., 2019; Calandriello et al., 2014), empirical examples of this
phenomenon are well-established (Sutton et al., 2011; Jaderberg et al., 2016).
Our work can be seen as a simple mathematical model for studying this
phenomenon in the Bayesian RL setting.

##  4 Background

In this section, we provide background on provably-efficient approaches to
addressing the exploration-exploitation trade-off in sequential decision-
making problems. The mechanisms employed by these approaches for representing
uncertainty will play a central role in our algorithm for learning a value-
based state abstraction.

###  4.1 Information-Directed Sampling

A central challenge that all sequential decision-making agents must confront
is that of exploration; an agent must strike a delicate balance between
acquiring new knowledge in the hope of improving future performance or
capitalizing on the information it has acquired thus far. Early results for
provably-efficient reinforcement-learning algorithms (designated PAC-MDP
(Strehl et al., 2009)) hinge on the sufficiency of an agent’s exploration
strategy for fully exploring the MDP (Kakade, 2003), typically based on a
principle of optimism in the face of uncertainty (Kearns and Singh, 2002;
Brafman and Tennenholtz, 2002; Bubeck and Cesa-Bianchi, 2012). In recent
years, exploration techniques that facilitate stronger theoretical guarantees
have come about by leveraging estimates of an agent’s epistemic uncertainty or
uncertainty stemming from parameter estimation (rather than the aleatoric
uncertainty driven by stochasticity in data) (Chapelle and Li, 2011; Russo and
Van Roy, 2016; Osband et al., 2016a; Agrawal and Jia, 2017; O’Donoghue et al.,
2018; Osband et al., 2019).

In the context of multi-armed bandit problems, an agent maintains uncertainty
over the individual reward or payoff functions at each arm. With uncertainty
in the rewards of all arms driving uncertainty over optimal actions, one
choice is for the agent to employ an exploration scheme based on Thompson
sampling (TS) (Thompson, 1933; Agrawal and Goyal, 2012, 2013; Russo and Van
Roy, 2016) whereby an agent acts optimally with respect to a single sample
drawn from its posterior beliefs at each time period. This idea naturally
scales to the full reinforcement-learning scenario wherein posterior beliefs
are maintained over the optimal action-value function $Q^{\star}$ (Osband et
al., 2016a; O’Donoghue et al., 2018; Osband et al., 2019).

A significant advance on the aforementioned exploration scheme is the
algorithmic design principle known as information-directed sampling (IDS)
(Russo and Van Roy, 2018). While previous approaches follow suit with Thompson
sampling and act optimally according to posterior samples, IDS algorithms
execute a policy at each time period that solves the following minimization
problem

$\displaystyle\pi_{h}$
$\displaystyle=\min\limits_{\pi\in\Delta(\mathcal{A})}\frac{(\mathbb{E}[\Delta_{h}(\pi)])^{2}}{\mathcal{I}(\theta;(S_{h},A_{h})|\mathcal{E}_{h-1})}$

where $\mathbb{E}[\Delta_{h}(\pi)]$ denotes the expected regret of policy
$\pi$ under the agent’s current posterior beliefs and
$\mathcal{I}(\theta;(S_{h},A_{h})|\mathcal{E}_{h-1})$ denotes the expected
information gain between the behavior at timestep $h$ and the environment
parameters $\theta$, conditioned on the history of episodes collected thus
far, $\mathcal{E}_{h-1}$. For a multi-armed bandit problem, $\theta$ reflects
the reward or payoffs at each arm whereas, for a MDP, $\theta$ captures the
environment transition function and reward function. IDS embodies an intuitive
principle that, rather than being biased exclusively towards the optimal
action of one posterior sample, an agent should be incentivized to take one or
more suboptimal actions so long as they are informative and revelatory of the
underlying environment, $\theta$. Here, the information gain term in the
denominator of the information ratio above quantifies this level of
informativity, weighing it against the agent’s desire to minimize regret over
its lifetime. Again, while simple and powerful, IDS is only a design principle
to guide the development of practical, efficient algorithms. In the next
section, we discuss a deliberate choice of how to represent an agent’s
posterior beliefs that yields a concrete instantiation of IDS for our
algorithm.

###  4.2 Hypermodels

When faced with a need to represent epistemic uncertainty, scalability becomes
an immediate challenge. For the popular choice of neural networks as
expressive function approximators, the common practice has been to employ
finite ensembles (Osband et al., 2016a; Lu and Van Roy, 2017) which maintain
several copies of neural network weights for a single architecture. A sample
from this posterior of $K$ ensemble members is acquired by first sampling
$z\sim\text{Uniform}(K)$ and then running the ensemble member indexed by $z$.

Recognizing the poor scaling of computationally efficiency for ensemble
sampling with increasing $K$, Dwaracherla et al. (2020) introduce hypermodels
as a scalable alternative to representing epistemic uncertainty. In short,
hypermodels rely on a space of indexes $\mathcal{Z}$ as well as a
corresponding reference distribution $p_{z}\in\Delta(\mathcal{Z})$. For a
given choice of base model $f_{\theta}:\mathcal{X}\mapsto\mathcal{Y}$
parameterized by $\theta\in\Theta$, a hypermodel with parameters $\nu$,
$\mathfrak{H}_{\nu}:\mathcal{Z}\mapsto\Theta$ maps a single index to a
particular instantiation of base model. Thus, by sampling an index $z\sim
p_{z}(\cdot)$, the function $f_{\mathfrak{H}_{\nu}(z)}$ represents a sample
from an approximate posterior. Given a dataset $\mathcal{D}$, the loss
function $\mathcal{L}(\nu,\mathcal{D})$ to optimize hypermodel parameters
$\nu$ will vary depending on the base model and task. We defer the definition
of $\mathcal{L}(\nu,\mathcal{D})$ to the next section.

##  5 Approach

###  5.1 Approximate Posterior over State Abstractions

Our core contribution is an approach for jointly learning a state abstraction
$\phi$ without prior knowledge of $Q^{\star}$. To do this, we leverage our
assumption of $Q^{\star}$ as being low rank and we explicitly maintain two
separate hypermodels $\mathfrak{H}^{\phi}_{\nu},\mathfrak{H}^{\psi}_{\nu}$ as
approximate posterior distributions over $\phi^{\star}$ and $\psi^{\star}$
respectively. It is important to note that while
$\mathfrak{H}^{\phi}_{\nu}:\mathcal{Z}\mapsto\Phi$ is a standard hypermodel
(Dwaracherla et al., 2020) mapping indices to instances of $\phi$,
$\mathfrak{H}^{\psi}_{\nu}:\mathcal{Z}\times\Phi\mapsto\Psi$ is a conditional
hypermodel over possible functions $\psi$. The intuition here is that once an
agent samples from its posterior beliefs over $\phi^{\star}$, the
corresponding sample from its beliefs over abstract-state values must be
conditioned on the particular sample
$\phi\sim\mathfrak{H}^{\phi}_{\nu}(\cdot)$. More succinctly, an agent’s
posterior beliefs over the optimal value function is obtained by first
sampling $z\sim p_{z}(\cdot)$ and then
$\hat{Q}^{\star}\sim\langle\mathfrak{H}^{\phi}_{\nu}(z),\mathfrak{H}^{\psi}_{\nu}(z,\mathfrak{H}^{\phi}_{\nu}(z))\rangle$.

Recall that we have yet to define the objective function for optimizing the
hypermodels to represent an approximate posterior over $Q^{\star}$. Given a
minibatch of past experiences $\tilde{\mathcal{D}}$, we optimize the following
loss function $\mathcal{L}(\nu,\tilde{\mathcal{D}})$:

$\displaystyle R(z_{\phi},z_{\psi},\nu)$
$\displaystyle\triangleq\lambda||\mathfrak{H}^{\phi}_{\nu}(z_{\phi})-\mathfrak{H}^{\phi}_{\nu_{0}}(z_{\phi})||^{2}_{2}+\lambda||\mathfrak{H}^{\psi}_{\nu}(z_{\psi},\mathfrak{H}^{\phi}_{\nu}(z_{\phi}))-\mathfrak{H}^{\psi}_{\nu_{0}}(z_{\psi},\mathfrak{H}^{\phi}_{\nu_{0}}(z_{\phi}))||^{2}_{2}$
$\displaystyle\hat{Q}^{\star}_{z_{\phi},z_{\psi}}(s,a)$
$\displaystyle\triangleq\langle\mathfrak{H}^{\phi}_{\nu}(z_{\phi})(s),\mathfrak{H}^{\psi}_{\nu}(z_{\psi},\mathfrak{H}^{\phi}_{\nu}(z_{\phi}))(a)\rangle$
$\displaystyle\mathcal{L}(\nu,\tilde{\mathcal{D}})=\mathbb{E}_{\begin{subarray}{c}z_{\phi}\sim
p_{z}(\cdot)\\\ z_{\psi}\sim
p_{z}(\cdot)\end{subarray}}[\frac{1}{m}\sum\limits_{(s,a,r,s^{\prime}\eta_{\phi},\eta_{\psi})\in\tilde{\mathcal{D}}}$
$\displaystyle(r+\gamma\max\limits_{a^{\prime}\in\mathcal{A}}\hat{Q}^{\star}_{z_{\phi},z_{\psi}}(s^{\prime},a^{\prime})+\eta_{\phi}^{T}z_{\phi}+\eta_{\psi}^{T}z_{\psi}-\hat{Q}^{\star}_{z_{\phi},z_{\psi}}(s,a))^{2}+R(z_{\phi},z_{\psi},\nu)]$

where $\nu_{0}$ denotes the initial vector of hypermodel parameters, the
$\eta_{\phi},\eta_{\psi}$ terms denote random Gaussian perturbations of the
target values, and $\lambda$ is a regularization coefficient. This loss
function encapsulates a randomized least-squares value iteration (RLSVI)
(Osband et al., 2016b) approach to maintaining an approximate posterior
distribution over $Q^{\star}$.

###  5.2 Variance-IDS

In the previous section, we specify how hypermodels can be used to maintain
approximate posterior distributions over state abstraction and abstract state
values. The final outstanding component that must be specified is how these
approximate posterior distributions can be folded into an algorithm that
instantiates IDS as the core exploration strategy. To do this, we leverage
variance-IDS as introduced in Russo and Van Roy (2018) and as specified for
hypermodels in Dwaracherla et al. (2020). For clarity, we present the
derivation of variance-IDS.

Recall the definition of the mutual information between two random variables
$X,Y$:

$\displaystyle\mathcal{I}(X;Y)$
$\displaystyle=\mathbb{E}_{X}[D_{\text{KL}}(p(Y|X)||p(Y))]$

Noting the definition contains a Kullback-Leibler (KL) divergence term, we
also define Pinkser’s inequality (Pinsker, 1960)

$\displaystyle D_{\text{TV}}(p(X)||q(X))$
$\displaystyle\leq\sqrt{\frac{1}{2}D_{\text{KL}}(p(X)||q(X))}$

where $D_{\text{TV}}(p(X)||q(X))$ denotes the total variation distance between
distributions $p,q$. Lastly, we note that the total variation distance is an
integral probability metric (IPM) (Müller, 1997) which, for a random variable
$X$ with support $\mathcal{X}$, is defined as:

$\displaystyle D_{\text{TV}}(p(X)||q(X))$
$\displaystyle=\sup\limits_{\begin{subarray}{c}f:\mathcal{X}\mapsto\mathbb{R}\\\
||f||_{\infty}\leq
1\end{subarray}}\mathbb{E}_{p(X)}[f(X)]-\mathbb{E}_{q(X)}[f(X)]$

where the supremum is taken with respect to all witness functions
$f:\mathcal{X}\mapsto\mathbb{R}$ with infinity norm bounded by 1. Putting all
the pieces together, we recall that IDS balances regret minimization with the
selection of informative actions, where informativeness of agent behavior is
measured by the (conditional) mutual information††Although not denoted here
for clarity, this is mutual information term is conditioned on the entire
history of all past interactions up to this point. between the state-action
pair observed at timestep $h$ and the true environment parameters $\theta$,
$\mathcal{I}(\theta;S_{h},A_{h})$. Following Russo and Van Roy (2018), this
term can be lower bounded as follows:

$\displaystyle\mathcal{I}(\theta;S_{h},A_{h})$
$\displaystyle=\mathbb{E}_{\theta}[D_{KL}(p(S_{h},A_{h}|\theta)||p(S_{h},A_{h}))]$
$\displaystyle\geq
2\mathbb{E}_{\theta}[(D_{TV}(p(S_{h},A_{h}|\theta)||p(S_{h},A_{h})))^{2}]$
$\displaystyle=2\mathbb{E}_{\theta}[(\sup\limits_{\begin{subarray}{c}f:\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}\\\
||f||_{\infty}\leq
1\end{subarray}}(\mathbb{E}_{p(S_{h},A_{h}|\theta)}[f(S_{h},A_{h})]-\mathbb{E}_{p(S_{h},A_{h})}[f(S_{h},A_{h})])^{2}]$
$\displaystyle\geq
2\mathbb{E}_{\theta}[(\mathbb{E}[Q^{\star}(S_{h},A_{h})|\theta]-\mathbb{E}[Q^{\star}_{h}(S_{h},A_{h})])^{2}]$
$\displaystyle=2\mathbb{V}[\mathbb{E}[Q^{\star}(S_{h},A_{h})|\theta]]$

where the steps follow from the definition of mutual information, Pinsker’s
inequality, the IPM form of the total variation distance, the definition of
supremum, and the definition of variance. Crucially, and just as with the
original mutual information term, all of the expectations and variances above
are conditioned on the past history of interactions up to this timestep. Thus,
the above shows that the information gain at timestep $h$ is lower bounded by
the extent to which $Q^{\star}$-values produced by action $a$ vary under the
identity of the environment parameters $\theta$. Consequently, those actions
with high variance in $Q^{\star}$-values under the agent’s current posterior
beliefs are deemed to be the most informative.

###  5.3 Posterior State-Abstraction Sampling

We present the full algorithm for PS2 as Algorithm 1 with the explicit choice
of hypermodels for representing an agent’s beliefs. The algorithm proceeds
according to variance-IDS by first drawing $K$ samples from the agent’s
current posterior beliefs over $Q^{\star}$ and computing the requisite
quantites for variance-IDS, namely the expected squared regret and expected
variance for each action. The resulting experience collected from the
environment is then accumulated via experience replay (Lin, 1992; Mnih et al.,
2015) for incrementally updating the hypermodels.

Data: Reference distributions $p_{z_{\phi}},p_{z_{\psi}}=\mathcal{N}(0,I)$,
Minibatch size $m$, Learning rate $\alpha$, Regularization parameters
$\lambda$, Noise variances $\sigma^{2}_{\phi}$, $\sigma^{2}_{\psi}$

Initialize $\mathfrak{H}^{\phi}_{\nu},\mathfrak{H}^{\psi}_{\nu}$

$\mathcal{D}\leftarrow\emptyset$

for  _t $1,\ldots,T$_ do

for  _h $1,\ldots,H$_ do

Sample
$Z=\\{(z_{\phi},z_{\psi})_{1},\ldots,(z_{\phi},z_{\psi})_{K}\\},z_{\phi}\sim
p_{z_{\phi}}(\cdot),z_{\psi}\sim p_{z_{\psi}}(\cdot)$

$\hat{Q}^{\star}_{(z_{\phi},z_{\psi})}(s,a)\triangleq\langle\mathfrak{H}^{\phi}_{\nu}(z_{\phi})(s),\mathfrak{H}^{\psi}_{\nu}(z_{\psi},\mathfrak{H}^{\phi}_{\nu}(z_{\phi}))(a)\rangle$

$\hat{\Delta}_{a}(s)\triangleq\frac{1}{|Z|}\sum\limits_{i=1}^{K}\max\limits_{a^{\star}\in\mathcal{A}}\hat{Q}^{\star}_{(z_{\phi},z_{\psi})_{i}}(s,a^{\star})-\hat{Q}^{\star}_{(z_{\phi},z_{\psi})_{i}}(s,a)]$

$\tilde{Z}(a)=\\{(z_{\phi},z_{\psi})|(z_{\phi},z_{\psi})\in
Z,a=\operatorname*{arg\,max}\limits_{a^{\star}\in\mathcal{A}}\hat{Q}^{\star}_{(z_{\phi},z_{\psi})}(s,a^{\star})\\}$

$\hat{v}_{a}(s)\triangleq\sum\limits_{a^{\star}\in\mathcal{A}}\frac{|\tilde{Z}(a^{\star})|}{|Z|}(\frac{1}{|\tilde{Z}(a^{\star})|}\sum\limits_{(z_{\phi},z_{\psi})\in\tilde{Z}(a^{\star})}\hat{Q}^{\star}_{(z_{\phi},z_{\psi})}(s,a)-\frac{1}{|Z|}\sum\limits_{(z_{\phi},z_{\psi})\in
Z}\hat{Q}^{\star}_{(z_{\phi},z_{\psi})}(s,a))^{2}$

$\pi(\cdot|s_{h})=\min\limits_{\pi\in\Delta(\mathcal{A})}\frac{\mathbb{E}_{a\sim\pi}[(\hat{\Delta}_{a}(s))^{2}]}{\mathbb{E}_{a\sim\pi}[\hat{v}_{a}(s)]}$

Take action $a_{h}\sim\pi(\cdot|s_{h})$ and observe $r_{h},s_{h+1}$

Sample random perturbations
$\eta_{\phi}\sim\mathcal{N}(0,\sigma^{2}_{\phi}I),\eta_{\psi}\sim\mathcal{N}(0,\sigma^{2}_{\psi}I)$

$\mathcal{D}\rightarrow\mathcal{D}\cup\\{(s_{h},a_{h},r_{h},s_{h+1},\eta_{\phi},\eta_{\psi})\\}$

Sample random minibatch $\tilde{\mathcal{D}}\stackrel{{\scriptstyle
m}}{{\sim}}\text{Uniform}(\mathcal{D})$

$\nu\leftarrow\nu-\alpha\nabla_{\nu}\mathcal{L}(\nu,\tilde{\mathcal{D}})$

end for

end for

Algorithm 1 Posterior State-Abstraction Sampling with Variance-IDS

For the multi-task learning setting, the algorithm is nearly identical, with
the caveat that index samples must be drawn for each hypermodel over abstract-
state values (with one hypermodel per task); the hypermodel over state
abstractions is shared across all tasks.

##  6 Experiments & Discussion

We recall that the primary goal of this paper is twofold, (1) offering PS2 as
a practical approach for synthesizing a state abstraction based on structure
within $Q^{\star}$ (as opposed to, for instance, structure in the transition
function) and (2) identifying when the pursuit of such latent structure is
particularly advantageous. In this section, we outline a recipe for randomly
generating contextual bandit problems that allow us to empirically address
these goals.

Concretely, given specific values for the number of states $(S)$, actions
$(A)$, and abstract states $(M)$, we generate contextual bandit problems by
randomly sampling a reward function with sparse latent structure. The
factorization $\mathbf{Q}^{\star}=\mathbf{\Phi}\mathbf{\Psi}^{T}$ is formed by
first generating $\mathbf{\Phi}\in\mathbb{R}^{S\times M}$ through the random
sampling of $S$ one-hot vectors of length $M$. The abstract-state values
$\mathbf{\Psi}\in\mathbb{R}^{A\times M}$ are then each drawn uniformly at
random from $[0,1]$. For our multi-task experiments, the procedure for
sampling $\mathbf{\Phi}$ remains unchanged and the procedure for sampling
$\mathbf{\Psi}_{t}$ is repeated for each task $t$. We define instantaneous
regret as
$\Delta_{t}=|\max\limits_{a^{\star}}Q^{\star}(s_{t},a^{\star})-Q^{\star}(s_{t},a_{t})|$
and define cumulative regret over $T$ episodes or time periods as
$\sum\limits_{t=1}^{T}\Delta_{t}$. For multi-task experiments, instantaneous
regret is summed across all tasks. In all experiments we use the Adam
optimizer (Kingma and Ba, 2014) with a learning rate of $0.001$, a batch size
of $1024$, $128$ index samples per timestep, noise variances of $0.25$, and
regularization parameter of $0.001$. All shading in figures denote $95\%$
confidence intervals computed across five random seeds. All hypermodels are
parameterized as Gaussian distributions with index samples of appropriate
dimension drawn iid from $\mathcal{N}(0,1)$. For each contextual-bandit
problem, we evaluate the following algorithms:

  * •

PS2-IDS – Algorithm 1 with two hypermodels over state abstractions and
abstract-state values respectively. The former hypermodel is parameterized as
a Gaussian distribution with diagonal covariance matrix. The latter hypermodel
over values is also represented as a Gaussian distribution with diagonal
covariance whose mean is a linear function of an input state abstraction
$\mathbf{\Phi}\in\mathbb{R}^{S\times M}$.

  * •

PS2-TS – Identical parameterization to PS2-IDS except, instead of using
variance-IDS, applies Thompson sampling for action selection by drawing a
single $Q^{\star}$ and then acting greedily with respect to the sample.

  * •

NoStateAbstraction – An implementation of variance-IDS that, instead of
learning a state abstraction, directly maintains a single hypermodel of the
agent’s beliefs over $\mathbf{Q}^{\star}\in\mathbb{R}^{S\times A}$.

  * •

TrueStateAbstraction – An implementation of variance-IDS that is given the
true state abstraction $\mathbf{\Phi}$ a priori and only maintains a
hypermodel for learning the corresponding abstract-state values.

  * •

Independent – Forr multi-task experiments with $T$ distinct tasks, this
algorithm maintains $T$ instances of PS2-IDS.

  * •

Random – Selects actions at each timestep uniformly at random.

(a) Single-task setting

(b) Multi-task setting with 10 tasks.

Figure 1: Contextual bandit with $S=10,A=10,M=5$.

(a) Single-task setting

(b) Multi-task setting with 10 tasks.

Figure 2: Contextual bandit with $S=30,A=30,M=5$.

Figures 1 and 2 show results for randomly generated contextual bandit problems
of two different sizes. Firstly, in the single-task setting, the relationship
between the algorithm that has the true state abstraction computed a priori
vs. the algorithm that does not pursue any abstraction at all reaffirms
existing empirical results that highlight the advantages of leveraging
structure in the value function (Abel et al., 2016). Notably, both PS2 methods
are able to achieve performance between these two extremes, identifying the
underlying abstraction of the environment to more efficiently arrive at
optimal behavior. This point becomes even more apparent in the multi-task
setting (Figures 1(b) and 2(b)) where interaction with multiple tasks
amplifies the signal provided to the agent for distilling the underlying state
abstraction. Due to the small problem size, Figure 1(b) shows little
improvement between PS2 and the algorithm that attempts to solve each task in
isolation. With a slightly larger problem in Figure 2(b), however, we observe
a substantial improvement in PS2 as it is able to better exploit information
from all tasks to capture shared structure. Finally, it is theoretically know
that IDS has a stronger performance guarantee than Thompson sampling (Russo
and Van Roy, 2018). Our experiments confirm this relationship with PS2-IDS
matching or outperforming Thompson sampling.

##  7 Conclusion

We have examined state abstractions as a mechanism for facilitating sample-
efficient reinforcement learning. While various forms of state abstraction
model structure in different components of a MDP, this work places particular
focus on those which attempt to leverage structure in the optimal value
function. While prior works that study this abstraction type have often been
unrealistic in their acquisition of the abstraction itself, we take a Bayesian
perspective and leverage insights from past work on provably-efficient, deep
exploration. Concretely, we introduce an algorithmic design principle wherein
an agent’s beliefs over the optimal value function factor into separate
posterior distributions over abstract states and abstract-state values,
respectively. Empirically, we demonstrate that an agent whose beliefs have
been factored in this manner can learn more efficiently than those that
attempt to directly estimate values and ignore latent structure altogether.
The experiments in this work have been limited to the contextual bandit
setting; identifying the right sampling procedure for generating random MDPs
with the appropriate structural properties that are conducive for any
algorithm, including PS2, to learn state abstractions is an active direction
for future work.

## References

  * Abel et al. [2016] David Abel, D Ellis Hershkowitz, and Michael L Littman.  Near optimal behavior via approximate state abstraction.  In _Proceedings of the 33rd International Conference on International Conference on Machine Learning-Volume 48_ , pages 2915–2923, 2016\.
  * Abel et al. [2018] David Abel, Dilip Arumugam, Lucas Lehnert, and Michael Littman.  State abstractions for lifelong reinforcement learning.  In _International Conference on Machine Learning_ , pages 10–19, 2018\.
  * Abel et al. [2019] David Abel, Dilip Arumugam, Kavosh Asadi, Yuu Jinnai, Michael L Littman, and Lawson LS Wong.  State abstraction as compression in apprenticeship learning.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 3134–3142, 2019.
  * Agarwal et al. [2020] Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun.  Flambe: Structural complexity and representation learning of low rank mdps.  _arXiv preprint arXiv:2006.10814_ , 2020.
  * Agrawal and Goyal [2012] Shipra Agrawal and Navin Goyal.  Analysis of thompson sampling for the multi-armed bandit problem.  In _Conference on learning theory_ , pages 39–1, 2012.
  * Agrawal and Goyal [2013] Shipra Agrawal and Navin Goyal.  Further optimal regret bounds for thompson sampling.  In _Artificial intelligence and statistics_ , pages 99–107, 2013\.
  * Agrawal and Jia [2017] Shipra Agrawal and Randy Jia.  Optimistic posterior sampling for reinforcement learning: worst-case regret bounds.  In _Advances in Neural Information Processing Systems_ , pages 1184–1194, 2017.
  * Auer et al. [2009] Peter Auer, Thomas Jaksch, and Ronald Ortner.  Near-optimal regret bounds for reinforcement learning.  In _Advances in neural information processing systems_ , pages 89–96, 2009.
  * Barreto et al. [2016] André MS Barreto, Rafael L Beirigo, Joelle Pineau, and Doina Precup.  Incremental stochastic factorization for online reinforcement learning.  In _Thirtieth AAAI Conference on Artificial Intelligence_ , 2016.
  * Baxter [1997] Jonathan Baxter.  A bayesian/information theoretic model of learning to learn via multiple task sampling.  _Machine learning_ , 28(1):7–39, 1997.
  * Baxter [2000] Jonathan Baxter.  A model of inductive bias learning.  _Journal of artificial intelligence research_ , 12:149–198, 2000.
  * Behzadian and Petrik [2018] Bahram Behzadian and Marek Petrik.  Feature selection by singular value decomposition for reinforcement learning.  In _Proceedings of the ICML Prediction and Generative Modeling Workshop_ , 2018.
  * Bellman [1957] Richard Bellman.  A markovian decision process.  _Journal of mathematics and mechanics_ , pages 679–684, 1957.
  * Bertsekas et al. [1988] Dimitri P Bertsekas, David A Castanon, et al.  Adaptive aggregation methods for infinite horizon dynamic programming.  1988\.
  * Brafman and Tennenholtz [2002] Ronen I Brafman and Moshe Tennenholtz.  R-max-a general polynomial time algorithm for near-optimal reinforcement learning.  _Journal of Machine Learning Research_ , 3(Oct):213–231, 2002.
  * Bubeck and Cesa-Bianchi [2012] Sébastien Bubeck and Nicolo Cesa-Bianchi.  Regret analysis of stochastic and nonstochastic multi-armed bandit problems.  _arXiv preprint arXiv:1204.5721_ , 2012.
  * Calandriello et al. [2014] Daniele Calandriello, Alessandro Lazaric, and Marcello Restelli.  Sparse multi-task reinforcement learning.  In _Advances in Neural Information Processing Systems_ , pages 819–827, 2014.
  * Candès and Plan [2010] Emmanuel J Candès and Yaniv Plan.  Matrix completion with noise.  _Proceedings of the IEEE_ , 98(6):925–936, 2010\.
  * Candès and Recht [2009] Emmanuel J Candès and Benjamin Recht.  Exact matrix completion via convex optimization.  _Foundations of Computational mathematics_ , 9(6):717, 2009.
  * Candès and Tao [2010] Emmanuel J Candès and Terence Tao.  The power of convex relaxation: Near-optimal matrix completion.  _IEEE Transactions on Information Theory_ , 56(5):2053–2080, 2010.
  * Caruana [1997] Rich Caruana.  Multitask learning.  _Machine learning_ , 28(1):41–75, 1997.
  * Chapelle and Li [2011] Olivier Chapelle and Lihong Li.  An empirical evaluation of thompson sampling.  In _Advances in neural information processing systems_ , pages 2249–2257, 2011.
  * Chen and Wainwright [2015] Yudong Chen and Martin J Wainwright.  Fast low-rank estimation by projected gradient descent: General statistical and algorithmic guarantees.  _arXiv preprint arXiv:1509.03025_ , 2015.
  * Chistov and Grigor’Ev [1984] Alexander L Chistov and D Yu Grigor’Ev.  Complexity of quantifier elimination in the theory of algebraically closed fields.  In _International Symposium on Mathematical Foundations of Computer Science_ , pages 17–31. Springer, 1984.
  * Dean and Givan [1997] Thomas Dean and Robert Givan.  Model minimization in markov decision processes.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , pages 106–111. AAAI Press, 1997.
  * D’Eramo et al. [2019] Carlo D’Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters.  Sharing knowledge in multi-task deep reinforcement learning.  In _International Conference on Learning Representations_ , 2019.
  * Dong et al. [2019] Shi Dong, Benjamin Van Roy, and Zhengyuan Zhou.  Provably efficient reinforcement learning with aggregated states.  _arXiv preprint arXiv:1912.06366_ , 2019.
  * Donoho and Stodden [2004] David Donoho and Victoria Stodden.  When does non-negative matrix factorization give a correct decomposition into parts?  In _Advances in neural information processing systems_ , pages 1141–1148, 2004.
  * Du et al. [2019] Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford.  Provably efficient rl with rich observations via latent state decoding.  In _International Conference on Machine Learning_ , pages 1665–1674, 2019.
  * Duan et al. [2019] Yaqi Duan, Tracy Ke, and Mengdi Wang.  State aggregation learning from markov transition data.  In _Advances in Neural Information Processing Systems_ , pages 4488–4497, 2019.
  * Dwaracherla et al. [2020] Vikranth Dwaracherla, Xiuyuan Lu, Morteza Ibrahimi, Ian Osband, Zheng Wen, and Benjamin Van Roy.  Hypermodels for exploration.  In _International Conference on Learning Representations_ , 2020.
  * Ferns et al. [2004] Norm Ferns, Prakash Panangaden, and Doina Precup.  Metrics for finite markov decision processes.  2004\.
  * Ferns et al. [2012] Norman Ferns, Pablo Samuel Castro, Doina Precup, and Prakash Panangaden.  Methods for computing state similarity in markov decision processes.  _arXiv preprint arXiv:1206.6836_ , 2012.
  * Gunasekar et al. [2013] Suriya Gunasekar, Ayan Acharya, Neeraj Gaur, and Joydeep Ghosh.  Noisy matrix completion using alternating minimization.  In _Joint European Conference on Machine Learning and Knowledge Discovery in Databases_ , pages 194–209. Springer, 2013.
  * Hardt [2014] Moritz Hardt.  Understanding alternating minimization for matrix completion.  In _2014 IEEE 55th Annual Symposium on Foundations of Computer Science_ , pages 651–660. IEEE, 2014.
  * Hardt et al. [2014] Moritz Hardt, Raghu Meka, Prasad Raghavendra, and Benjamin Weitz.  Computational limits for matrix completion.  In _Conference on Learning Theory_ , pages 703–725, 2014.
  * Hazan et al. [2019] Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest.  Provably efficient maximum entropy exploration.  In _International Conference on Machine Learning_ , pages 2681–2691, 2019.
  * Jaderberg et al. [2016] Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu.  Reinforcement learning with unsupervised auxiliary tasks.  _arXiv preprint arXiv:1611.05397_ , 2016.
  * Jain et al. [2013] Prateek Jain, Praneeth Netrapalli, and Sujay Sanghavi.  Low-rank matrix completion using alternating minimization.  In _Proceedings of the forty-fifth annual ACM symposium on Theory of computing_ , pages 665–674, 2013.
  * Jiang et al. [2015] Nan Jiang, Alex Kulesza, and Satinder Singh.  Abstraction selection in model-based reinforcement learning.  In _International Conference on Machine Learning_ , pages 179–188, 2015.
  * Jin et al. [2020] Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu.  Reward-free exploration for reinforcement learning.  _arXiv preprint arXiv:2002.02794_ , 2020.
  * Jong and Stone [2005] Nicholas K Jong and Peter Stone.  State abstraction discovery from irrelevant state variables.  In _Proceedings of the 19th international joint conference on Artificial intelligence_ , pages 752–757, 2005.
  * Kakade [2003] Sham Machandranath Kakade.  _On the sample complexity of reinforcement learning_.  PhD thesis, 2003.
  * Kearns and Singh [2002] Michael Kearns and Satinder Singh.  Near-optimal reinforcement learning in polynomial time.  _Machine learning_ , 49(2-3):209–232, 2002.
  * Keshavan et al. [2010a] Raghunandan H Keshavan, Andrea Montanari, and Sewoong Oh.  Matrix completion from a few entries.  _IEEE transactions on information theory_ , 56(6):2980–2998, 2010a.
  * Keshavan et al. [2010b] Raghunandan H Keshavan, Andrea Montanari, and Sewoong Oh.  Matrix completion from noisy entries.  _Journal of Machine Learning Research_ , 11(Jul):2057–2078, 2010b.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Konidaris [2019] George Konidaris.  On the necessity of abstraction.  _Current opinion in behavioral sciences_ , 29:1–7, 2019\.
  * Koren et al. [2009] Yehuda Koren, Robert Bell, and Chris Volinsky.  Matrix factorization techniques for recommender systems.  _Computer_ , 42(8):30–37, 2009.
  * Krishnamurthy and Singh [2013] Akshay Krishnamurthy and Aarti Singh.  Low-rank matrix and tensor completion via adaptive sampling.  In _Advances in neural information processing systems_ , pages 836–844, 2013.
  * Krishnamurthy and Singh [2014] Akshay Krishnamurthy and Aarti Singh.  On the power of adaptivity in matrix completion and approximation.  _arXiv preprint arXiv:1407.3619_ , 2014.
  * Lake et al. [2017] Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel J Gershman.  Building machines that learn and think like people.  _Behavioral and brain sciences_ , 40, 2017.
  * Li et al. [2006] Lihong Li, Thomas J. Walsh, and Michael L. Littman.  Towards a unified theory of state abstraction for MDPs.  In _Proceedings of the International Symposium on Artificial Intelligence and Mathematics_ , 2006.
  * Lin [1992] Long-Ji Lin.  Self-improving reactive agents based on reinforcement learning, planning and teaching.  _Machine learning_ , 8(3-4):293–321, 1992.
  * Lu and Van Roy [2017] Xiuyuan Lu and Benjamin Van Roy.  Ensemble sampling.  In _Advances in neural information processing systems_ , pages 3258–3266, 2017.
  * Misra et al. [2019] Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford.  Kinematic state abstraction and provably efficient rich-observation reinforcement learning.  _arXiv preprint arXiv:1911.05815_ , 2019.
  * Mnih et al. [2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  _nature_ , 518(7540):529–533, 2015.
  * Müller [1997] Alfred Müller.  Integral probability metrics and their generating classes of functions.  _Advances in Applied Probability_ , 29(2):429–443, 1997.
  * Osband and Van Roy [2017] Ian Osband and Benjamin Van Roy.  Why is posterior sampling better than optimism for reinforcement learning?  In _International Conference on Machine Learning_ , pages 2701–2710, 2017.
  * Osband et al. [2016a] Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy.  Deep exploration via bootstrapped dqn.  In _Advances in neural information processing systems_ , pages 4026–4034, 2016a.
  * Osband et al. [2016b] Ian Osband, Benjamin Van Roy, and Zheng Wen.  Generalization and exploration via randomized value functions.  In _International Conference on Machine Learning_ , pages 2377–2386, 2016b.
  * Osband et al. [2019] Ian Osband, Benjamin Van Roy, Daniel J Russo, and Zheng Wen.  Deep exploration via randomized value functions.  _Journal of Machine Learning Research_ , 20(124):1–62, 2019.
  * O’Donoghue et al. [2018] Brendan O’Donoghue, Ian Osband, Remi Munos, and Volodymyr Mnih.  The uncertainty bellman equation and exploration.  In _International Conference on Machine Learning_ , pages 3836–3845, 2018.
  * Pinsker [1960] Mark Semenovich Pinsker.  Information and information stability of random variables and processes.  1960\.
  * Puterman [1994] Martin L. Puterman.  _Markov Decision Processes—Discrete Stochastic Dynamic Programming_.  John Wiley & Sons, Inc., New York, NY, 1994.
  * Recht [2011] Benjamin Recht.  A simpler approach to matrix completion.  _Journal of Machine Learning Research_ , 12(Dec):3413–3430, 2011.
  * Recht and Ré [2013] Benjamin Recht and Christopher Ré.  Parallel stochastic gradient algorithms for large-scale matrix completion.  _Mathematical Programming Computation_ , 5(2):201–226, 2013.
  * Russo [2020] Daniel Russo.  Approximation benefits of policy gradient methods with aggregated states.  _arXiv preprint arXiv:2007.11684_ , 2020.
  * Russo and Van Roy [2016] Daniel Russo and Benjamin Van Roy.  An information-theoretic analysis of thompson sampling.  _The Journal of Machine Learning Research_ , 17(1):2442–2471, 2016.
  * Russo and Van Roy [2018] Daniel Russo and Benjamin Van Roy.  Learning to optimize via information-directed sampling.  _Operations Research_ , 66(1):230–252, 2018.
  * Shah et al. [2020] Devavrat Shah, Dogyoon Song, Zhi Xu, and Yuzhe Yang.  Sample efficient reinforcement learning via low-rank matrix estimation.  _arXiv preprint arXiv:2006.06135_ , 2020.
  * Shannon [1959] Claude E. Shannon.  Coding theorems for a discrete source with a fidelity criterion.  _IRE Nat. Conv. Rec., March 1959_ , 4:142–163, 1959.
  * Singh et al. [1995] Satinder P Singh, Tommi Jaakkola, and Michael I Jordan.  Reinforcement learning with soft state aggregation.  In _Advances in neural information processing systems_ , pages 361–368, 1995.
  * Strehl et al. [2009] Alexander L Strehl, Lihong Li, and Michael L Littman.  Reinforcement learning in finite mdps: Pac analysis.  _Journal of Machine Learning Research_ , 10(Nov):2413–2444, 2009.
  * Sutton [1988] Richard S Sutton.  Learning to predict by the methods of temporal differences.  _Machine learning_ , 3(1):9–44, 1988.
  * Sutton and Barto [1998] Richard S Sutton and Andrew G Barto.  Introduction to reinforcement learning.  1998\.
  * Sutton et al. [2000] Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour.  Policy gradient methods for reinforcement learning with function approximation.  In _Advances in neural information processing systems_ , pages 1057–1063, 2000.
  * Sutton et al. [2011] Richard S Sutton, Joseph Modayil, Michael Delp, Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup.  Horde: A scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction.  In _The 10th International Conference on Autonomous Agents and Multiagent Systems-Volume 2_ , pages 761–768, 2011.
  * Thompson [1933] William R Thompson.  On the likelihood that one unknown probability exceeds another in view of the evidence of two samples.  _Biometrika_ , 25(3/4):285–294, 1933.
  * Tsitsiklis and Van Roy [1996] John N Tsitsiklis and Benjamin Van Roy.  Feature-based methods for large scale dynamic programming.  _Machine Learning_ , 22(1-3):59–94, 1996.
  * Van Roy [2006] Benjamin Van Roy.  Performance loss bounds for approximate value iteration with state aggregation.  _Mathematics of Operations Research_ , 31(2):234–244, 2006.
  * Watkins and Dayan [1992] Christopher JCH Watkins and Peter Dayan.  Q-learning.  _Machine learning_ , 8(3-4):279–292, 1992.
  * Whitt [1978] Ward Whitt.  Approximations of dynamic programs, i.  _Mathematics of Operations Research_ , 3(3):231–243, 1978.
  * Williams [1992] Ronald J Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine learning_ , 8(3-4):229–256, 1992.
  * Yang and Wang [2019a] Lin Yang and Mengdi Wang.  Sample-optimal parametric q-learning using linearly additive features.  In _International Conference on Machine Learning_ , pages 6995–7004, 2019a.
  * Yang and Wang [2019b] Lin F Yang and Mengdi Wang.  Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound.  _arXiv preprint arXiv:1905.10389_ , 2019b.
  * Zhang and Wang [2019] Anru Zhang and Mengdi Wang.  Spectral state compression of markov processes.  _IEEE Transactions on Information Theory_ , 66(5):3202–3231, 2019.