# Distortion Audio Effects:
Learning How to Recover the Clean Signal

###### Abstract

Given the recent advances in music source separation and automatic mixing,
removing audio effects in music tracks is a meaningful step toward developing
an automated remixing system. This paper focuses on removing distortion audio
effects applied to guitar tracks in music production. We explore whether
effect removal can be solved by neural networks designed for source separation
and audio effect modeling.

Our approach proves particularly effective for effects that mix the processed
and clean signals. The models achieve better quality and significantly faster
inference compared to state-of-the-art solutions based on sparse optimization.
We demonstrate that the models are suitable not only for declipping but also
for other types of distortion effects. By discussing the results, we stress
the usefulness of multiple evaluation metrics to assess different aspects of
reconstruction in distortion effect removal.

##  1 Introduction

With the emergence of musical recordings, audio effects have become
indispensable in the music production process. They are used by musicians as a
creative tool to alter the sound of their instruments, and by sound engineers
to craft a balanced mix from multiple recording tracks [1].

For the task of mixing and automatic remixing [2], the dry (i.e., unprocessed)
source tracks are required. Given the recent advances in automatic mixing [3,
4] and music source separation (MSS) [5], a system could facilitate the
adjustment of a stereo mixture to the taste and preferences of the user
similar to [6]. However, when separating sources with a system trained on
music stems (e.g., `MUSDB18` [7]) the mixing process is not considered, and,
hence, the output of such a system contains the wet (i.e., processed) signal.
As nonlinear distortion is one of the most commonly used effects for electric
instruments, this work focuses on musical distortion effects that are used for
aesthetic means (e.g., guitar overdrive/distortion pedals) and applied in the
process of mixing (e.g., tape saturation). These distortion effects result in
added harmonics, intermodulation distortion, and a compressed sound [8].

This paper investigates different deep neural network (DNN) approaches
regarding their applicability to the audio effect removal problem. Our
contributions can be summarized as follows:

  * •

We show that recovering the clean signal from clipped or overdriven guitar
signals can be efficiently solved with neural networks designed for source
separation. The models achieve high quality and fast inference in contrast to
solutions based on sparse optimization.

  * •

We found that the superior performance of the models evaluated on the de-
overdrive task can be traced back to superimposing the overdriven signal with
the clean signal. We show that the architectures are suitable not only for
declipping but also for other types of distortion effects.

  * •

By discussing the results, we highlight that the metrics under evaluation
prove beneficial in measuring different aspects of the reconstruction and can
be advised for further investigations.

This work is organized as follows: Sec. 2 gives a formal introduction to audio
effect removal and introduces the types of distortions that were used
throughout this study. Sec. 3 briefly discusses previous work on iterative and
DNN-based declipping approaches. The methods under evaluation are outlined in
Sec. 4. Sec. 5 describes the data that were used for training, reports details
about the experimental setup, and presents the chosen objective evaluation
metrics. In Sec. 6, we evaluate the results of the comparative study of four
different neural network architectures on the task of distortion removal in
guitar signals for the SoX overdrive implementation. Then, we compare the same
architectures on the declipping task to one state-of-the-art declipping
algorithm using guitar signals as well as generic music signals. Lastly, Sec.
7 gives a conclusion and presents an outlook for future work. Audio examples
are available online at joimort.github.io/distortionremoval/.

##  2 Problem Formulation

We introduce audio effect removal as the task of recovering the original
discrete audio signal $\mathbf{x}\in\mathbb{R}^{n}$ from the processed
discrete signal $\mathbf{y}\in\mathbb{R}^{n}$, which is obtained by applying
the possibly nonlinear and time-varying function $f$ to the signal
$\mathbf{x}$:

$\begin{split}\mathbf{y}&=g(\mathbf{x})=\alpha
f(\mathbf{x})+(1-\alpha)\mathbf{x},\end{split}$ (1)

with $\alpha\in\left[0,1\right]$ denoting the weight of the wet signal, and
$g$ the summation function of the dry and wet signal. The goal is to find an
estimation of the original signal $\mathbf{\hat{x}}$ by estimating the inverse
function $\mathbf{\hat{x}}=\hat{g}^{-1}(\mathbf{y})$.

Generally, distortion effects clip the input signal and can be divided into
systems that apply hard-clipping or soft-clipping.

Figure 1: Example of different distortion types applied to a guitar signal.
The input signal $x(k)$ is amplified before being modified by a wave-shaper
function.

As this work focuses on both types of distortion, we introduce the following
generic formulation of a wave-shaper that maps the amplified signal
$\mathbf{x}_{\gamma}=10^{\frac{\gamma}{20}}\mathbf{x}$ to a fixed range:

$f(\mathbf{x})=\lambda(\mathbf{x}_{\gamma})\quad\textrm{with}\quad\lambda:\mathbb{R}\rightarrow[-\theta_{\mathrm{c}},\theta_{\mathrm{c}}].$
(2)

Here, $\gamma$ denotes the gain in decibels, $\lambda$ the arbitrary wave-
shaper function, and $\theta_{\mathrm{c}}$ the fixed clipping threshold. For
the case of hard-clipping, $\lambda$ is defined as:

$\lambda_{\mathrm{hc}}(x_{\gamma,k})=\begin{cases}x_{\gamma,k},&\text{if
}|x_{\gamma,k}|\leq\theta_{\mathrm{c}}\\\
\theta_{\mathrm{c}}\mathrm{sgn}(x_{\gamma,k})&\text{otherwise},\end{cases}$
(3)

with $\mathrm{sgn}$ denoting the sign function, and $x_{\gamma,k}$ the $k$-th
time sample of the amplified signal. Fig. 1 highlights the difference between
different distortion types. Hard-clipping cuts off the amplitude when it
exceeds a defined threshold (as typical for saturation in digital signal
processing). Soft-clipping (e.g.,
$\lambda_{\mathrm{tanh}}(x_{\gamma,k})\\!=\\!\mathrm{tanh}(x_{\gamma,k})$)
gradually applies a smooth transition before reaching a fully saturated state
(as typical for saturation in analog amplifiers).

Modeling the characteristics of distortion pedals in reality is more complex:
[9] provides an overview on different methods and discusses DNN-based
approaches. In order to simplify the problem for our investigation, we focus
on wave-shaping. The overdrive algorithm of the audio editing software SoX
[10] serves as an example of soft-clipping ($\lambda_{\mathrm{sox}}$) but, in
contrast to (2), it is dependent on previous samples. Furthermore, it blends
the wet signal with the dry one ($\alpha<1$).

##  3 RELATED WORK

To the best of our knowledge, there has been no previous research on
distortion audio effect removal. Therefore, this section outlines the most
relevant iterative and DNN-based declipping approaches since declipping is a
special case of distortion audio effect removal.

###  3.1 Iterative Declipping Methods

Previous research on approaches to declipping has focused mainly on
unsupervised algorithms that recover the signal under the assumption of a
generic regularization such as signal sparsity [11]. Usually, these approaches
target the hard-clipping case only (see (3)). While early approaches were
based on auto-regressive models, (e.g., [12]) recent state-of-the-art methods
evolved by combining ideas from inverse problems and sparse regularization
[13].

Recently, [11] and [13] discussed popular declipping algorithms. One of the
current state-of-the-art methods is A-SPADE, which will serve as a baseline
for this study. The algorithm is briefly introduced in Sec. 4.

###  3.2 DNN-Based Declipping

In contrast to iterative algorithms, to date, there are only few contributions
comprising supervised DNNs.

Kashani et al. [14] introduced a declipping method based on the U-Net
architecture [15]. It operates on magnitude spectrograms while the waveform of
the output is obtained by reusing the phase information from the distorted
input signal. The system is trained and evaluated on pairs of hard-clipped and
clean speech samples.

Mack and Habets [16] proposed an architecture comprising a BLSTM-based deep
filtering method that works on complex spectrograms and hence also considers
phase information. Similar to [14], they train the system on speech data only.
Unlike any other approach, they not only investigate the system on the hard-
clipping case but also on the soft-clipping case.

Tanaka et al. recently proposed APPLADE [17], a declipping method that takes
advantage of the sparse optimization techniques described above together with
deep learning. Accordingly, they embed a DNN in the iterative algorithm to
enhance the thresholding operation. They report a slightly higher performance
than previous algorithms, better robustness to mismatches between training and
test data, and faster inference.

##  4 METHODS

This section describes four neural network architectures that we selected from
the literature and evaluated on the distortion removal task. We approach the
distortion effect removal problem from the perspective of filtering the added
harmonics and intermodulation distortion, similar to [16]. Therefore, we
include one model from the domain of audio effect modeling and three
architectures originally proposed for music source separation.111Two of the
methods under evaluation, Demucs and Wave-U-Net, do not explicitly filter the
signals in the audio domain, rather they perform a nonlinear mapping.
Nevertheless, they were both successfully employed for MSS, which is a
filtering problem. For the latter models, instead of multiple output channels
for different sources, we use only one output channel and consider them as
general audio-to-audio transformation architectures.

CRAFx was proposed as a system for modeling time-varying audio effects with a
neural network [18, 9]. The end-to-end model operates on the signal in the
time domain and is divided into an adaptive front-end (encoder), a bi-
directional long-short-term-memory (BLSTM)-based structure that applies the
modeled effect in the latent space, and a synthesis back-end (decoder). In
contrast to the other DNN-based models presented in this section, this
architecture employs architectural priors (e.g., learnable nonlinear
activations) in the context of audio effects.

Open-Unmix (UMX) was introduced as a reference implementation for music source
separation [19]. The architecture is based on the BLSTM model from [20] and
uses magnitude spectrograms as input features. The essential element of Open-
Unmix is its three-layer BLSTM network that enables to learn both long- and
short-time dependencies [21]. An element-wise multiplication of the input
spectrograms with the estimated masks yields the final output. Commonly,
spectrogram-based source separation models are compared with the oracle
performance of an ideal ratio mask (IRM) that is defined as the ratio between
the reference and the test spectrogram [22] in decibels. For reconstruction,
the phase of the input signal is used. The model was adapted for a sampling
rate of $f_{s}=$16\text{\,}\mathrm{kHz}$$. We include this model in our
evaluation as a standard frequency domain MSS model that relates to the BLSTM-
based declipping model from [16] (cf. Sec. 3.2).

Wave-U-Net was proposed as one of the first end-to-end approaches for music
source separation based on time domain signals [23]. Hence, it incorporates
not only the magnitude but also the phase of music signals. It adapts the
U-Net architecture [15] to one-dimensional audio signals. We decreased the
models’ number of learnable parameters from $17$M to approximately $1$M by
reducing the number of layers from $12$ to $8$ resulting in a reduced
receptive field, as we experienced overfitting for our datasets. Since the
model was successfully employed not only for source separation but automatic
mixing as well [3], we assume a general suitability for audio-to-audio
transformation tasks. Moreover, the model is highly related to the U-Net-based
declipping model from [14] (cf. Sec. 3.2).

Demucs was initially designed to be an end-to-end model for music source
separation in the time domain [24]. While it builds on the Wave-U-Net model,
it introduces several improvements to the architecture. The model comprises a
convolutional encoder, a BLSTM structure, and a convolutional decoder. As for
Wave-U-Net, we decreased the number of trainable parameters from $66$M to $1$M
by reducing the number of blocks from $12$ to $6$ resulting in a reduced
receptive field.

A-SPADE was introduced as a sparsity-based declipping algorithm that
outperforms previous similar approaches [25]. For each time frame of the
clipped signal $\mathbf{y}$, it approximates a solution of the following
problem:

$\min_{\mathbf{x,z}}||\mathbf{z}||_{0}\mkern 9.0mu\textrm{s.t.}\mkern
9.0mu\mathbf{x}\in\Gamma(\mathbf{y})\mkern 9.0mu\textrm{and}\mkern
9.0mu||\mathcal{F}(\mathbf{x})-\mathbf{z}||_{2}\leq\epsilon,$ (4)

where $\mathbf{z}$ denotes the unknown discrete Fourier coefficients of each
time frame and $\mathcal{F}$ the Fourier transform operator. $\Gamma$ is
defined as the feasible space of solutions (i.e., clipping consistency
constraint). We included the algorithm in the evaluation of the declipping
task as a baseline that delivers state-of-the-art performance [11, 13].

##  5 EXPERIMENTS

Our experiments focus on the following three scenarios: a) We conducted
experiments on guitar recordings that were processed using the overdrive
algorithm of the audio editing software SoX [10] (CEG-OD). b) We performed the
same experiments as in the previous scenario while processing the same clean
guitar recordings with hard-clipping (CEG-HC). c) We tested the systems on an
extensive dataset comprising various hard-clipped sounds (SignalTrain-HC) to
evaluate their performance against a popular declipping algorithm when the
models are trained given an increase in the amount and variety of data.

###  5.1 Data

The models were trained on two different datasets to assess the distortion
audio effect removal capabilities. The audio signals from a dataset that
contains a single instrument class (e.g., electric guitar) exhibit consistent
signal statistics. In order to restrict the statistics that need to be modeled
in a first step, we chose to concentrate on dry guitar samples as target data.

Since a large-scale, polyphonic dataset from clean electric guitar sounds is,
to the best of our knowledge, not available222A publicly available guitar
dataset for the recognition of audio effects exists (IDMT-SMT-Audio-Effects
[26]). However, the dataset contains primarily homogeneous monophonic sounds
and therefore, we choose to use the CEG dataset instead., we used an internal
dataset, which we refer to as CEG (Clean Electric Guitar) dataset. The
monaural data were gathered from various sources, mainly commercial audio loop
packages and recordings of solo guitar, and has a duration of
$1.68\text{\,}\mathrm{h}$. All signals were re-sampled to a common sampling
rate of $16\text{\,}\mathrm{kHz}$ in order to speed up convergence during
training. To create the input dataset CEG-OD, the overdrive algorithm of SoX
was applied to the data using five uniformly sampled gain levels in the range
of $\gamma\in[20,50]$$\mathrm{dB}$. Likewise, the input dataset for the hard-
clipping task, CEG-HC, was created using hard-clipping (see (3)) with gain
levels from the same distribution. Both datasets have a total length of
$8.4\text{\,}\mathrm{h}$.

Although CEG represents a good source of data for our experiments due to its
specificity of timbre, it remains a limited resource in terms of size and
variety. Before attempting to train a system to handle recordings in real
environments (e.g., a commercial song), we need to investigate how the current
models at our disposal handle the availability of more and diverse training
data. For this purpose, we also performed experiments on the SignalTrain
dataset, which consists of more than $24\text{\,}\mathrm{h}$ of music and
randomly-generated test signals [27]. By applying hard-clipping to these clean
data using a uniformly-sampled input SDR value in the range
$\mathrm{SDR_{inp}}\in[1,20]$\mathrm{dB}$$, we created SignalTrain-HC. During
evaluation, we applied each input SDR in the set
$\mathrm{SDR_{inp}}\in\\{1,3,5,7,10,15,20\\}$\mathrm{dB}$$ to each sample in
the test set.

Each dataset was split into non-overlapping subsets for training (80%),
validation (10%) and testing (10%). We evaluated the models on the test split.

###  5.2 Experimental Setup

During the supervised training procedure, Adam [28] was used as optimizer with
initial learning rates according to the model’s original implementations. The
learning rate was reduced by a factor of 10 after 150 epochs of no decrease in
the validation loss. All models were optimized using the source-to-distortion
ratio (SDR) between their full output sequences $\mathbf{\hat{x}}$ and the
respective target sequences $\mathbf{x}$ in each batch $\mathcal{B}$ of $N$
elements (not to be confused with the definition in the `BSS_eval` toolkit
[29]):

$\mathcal{L}_{\mathcal{B}}(\mathbf{x},\hat{\mathbf{x}})=\frac{1}{N}\sum_{i\in\mathcal{B}}10\log_{10}\left(||\mathbf{x}_{i}||^{2}/||\mathbf{x}_{i}-\hat{\mathbf{x}_{i}}||^{2}\right).$
(5)

We stopped all trainings after 1000 epochs. All models processed audio
sequences that are randomly extracted from each clip in the dataset; the
length of the extracted sequences is equal to $2\text{\,}\mathrm{s}$
($2.3\text{\,}\mathrm{s}$ for CRAFx due to its architecture). We used a batch
size of 16 for all experiments.

###  5.3 Objective Metrics

In speech enhancement and source separation, the ubiquitous measure to
estimate the quality of a system is the SDR. However, applying (2) to a signal
does not retain its energy. Therefore, we follow the approach of [30]: the
scale-invariant SDR (SI-SDR) is obtained by rescaling the target signal $s$
such that the residual $s-\hat{s}$ is orthogonal to $s$ by using the optimal
scaling factor $\hat{s}^{T}s/||s||^{2}$.

An evaluation exclusively based on the objective similarity of the signals
does not necessarily imply a correlation with human perception [31, 32].
Accordingly, we observed that the SI-SDR scores occasionally disagreed with
our qualitative evaluation. Therefore, we also considered three metrics based
on human perception.

The perceptual evaluation of audio quality (PEAQ) [33] is a widely used
perceptual metric [11, 13, 34, 35] that measures the amount of degradation
between two audio signals. The output of PEAQ is an Overall Difference Grade
(ODG), which can reach values between $0$ (imperceptible impairment) and $-4$
(very annoying impairment). Even though PEAQ is used in declipping and audio
restoration studies [11, 13, 34], it was developed for audio codecs.

The R-nonlin metric [36], in contrast, was developed specifically for
detecting nonlinear distortions and, like PEAQ, considers the human auditory
system. $R_{\mathrm{nonlin}}$ is defined between 0 (high distortion) and 1 (no
distortion).

The Fréchet audio distance (FAD) was recently proposed as a reference-free
evaluation metric for music enhancement algorithms. It has shown to correlate
more with human perception than the SDR [37]. In order to obtain the FAD, the
embedding statistics of both the whole clean and distorted test set are
generated using a VGGish model [38]. The FAD is calculated based on the
Fréchet distance between two multivariate Gaussians computed from both the
test and the reference embeddings. [39].

##  6 RESULTS

Figure 2: Box plot of scores for the CEG-OD dataset. The boxes show the first
and third quartile of the data while the median is indicated with a line in
the box. Higher scores indicate superior performance except for FAD. The
results suggest that U-Net-based models performing convolutions in the time
domain are the most promising approach to solving the task of overdrive
removal.

In this section, we provide the results of the experiments introduced in the
previous section.

###  6.1 De-Overdrive (CEG-OD)

Fig. 2 shows the results of the models that remove overdrive from guitar
tracks.

Firstly, in SI-SDR, both Demucs and Wave-U-Net perform exceptionally well and
even outperform the ideal-ratio-mask by more than $24\text{\,}\mathrm{dB}$.
While CRAFx yields considerably worse performance, it also surpasses the IRM
oracle. UMX is the least performing of all the models in our comparison. It
should be noted that for UMX, a model that operates on magnitude spectrograms,
the IRM represents its upper limit in performance.

Similar results are obtained with PEAQ, $R_{\mathrm{nonlin}}$ and FAD: while
Demucs and Wave-U-Net yield the best scores and surpass the IRM, the ones for
CRAFx and UMX are considerably worse. It seems that directly processing the
signals in the time domain using a U-Net-based architecture represents the
most promising approach for the removal of the overdrive effect.

###  6.2 Declipping (CEG-HC)

Figure 3: Box plot of scores for the CEG-HC dataset. It is difficult to
clearly determine the best model, although Demucs represents a good compromise
for all the metrics.

Fig. 3 shows the results on the task of declipping on guitar recordings.
Generally, we experienced a drop in the scores: now the models have been
trained on declipping, which is an ill-posed problem, as missing parts of the
signal need to be reconstructed. Additionally, we report scores for our
declipping baseline, A-SPADE.

Despite the general performance drop, Demucs surpasses the results of A-SPADE
in terms of SI-SDR by almost $1\text{\,}\mathrm{dB}$. While UMX and Wave-U-Net
yield similar performances, CRAFx is the method with lowest scores.

Regarding PEAQ, no method surpasses the IRM and the A-SPADE algorithm. As
before, Demucs and Wave-U-Net have similar performance: while PEAQ slightly
favors the latter, SI-SDR favors the former. In contrast to the previous
results, CRAFx has no significant advantage over UMX.

While Demucs achieves the best score for $R_{\mathrm{nonlin}}$ among the
neural models, it does not surpass A-SPADE since the difference in their score
is marginal. Interestingly, UMX achieves better results than CRAFx and Wave-U-
Net. As the $R_{\mathrm{nonlin}}$ metric was explicitly designed to detect
nonlinear distortions, we conclude that UMX’ outputs contain fewer nonlinear
distortions than the outputs of CRAFx and Wave-U-Net.

Surprisingly, when computing the FAD, UMX outperforms all other methods,
including A-SPADE. This might be accounted to the fact that the FAD is based
on the mel spectrum, whereas UMX optimizes the magnitude spectrogram. While
Demucs and CRAFx outperform A-SPADE in FAD as well, the score for Wave-U-Net
is slightly worse. The FAD for the IRM is relatively small because the
difference between the reference and test embeddings from the VGGish model can
be traced back to the masking operation and quantization. Demucs seems to
constitute a good compromise regarding performance since it yields first- or
second-best results for all metrics.

###  6.3 Declipping (SignalTrain-HC)

Figure 4: Box plot of scores for the SignalTrain-HC dataset. Demucs can be
regarded as the best model regarding the median score across all metrics.

Fig. 4 shows the results for declipping SignalTrain-HC. Due to the lower gains
that are used to prepare the data (see Sec. 5.1), the overall performance
seems to be superior but cannot be directly compared to the previous results.
Because of its considerably worse performance in the previous task, CRAFx was
left out of the evaluation.

Demucs surpasses all other models in SI-SDR, including IRM and A-SPADE.
Moreover, Wave-U-Net and UMX both surpass A-SPADE but not the IRM. PEAQ gives
a similar ranking of the methods: only UMX cannot reach the A-SPADE baseline.
None of the neural methods surpasses the IRM. In terms of
$R_{\mathrm{nonlin}}$, the same ranking is obtained, with Demucs surpassing,
Wave-U-Net reaching, and UMX just missing A-SPADE. The FAD highlights that UMX
and Demucs deliver comparable performance, outperforming the baseline, but not
surpassing the IRM.

When only looking at SI-SDR or PEAQ, we notice the superiority of the time
domain models. Future research should investigate whether the waveform in the
time domain is the best input representation for the task, compared to, e.g.,
the real and imaginary part of a spectrogram [40] or both the waveform and the
spectrogram [41]. Ultimately, Demucs can be considered the best model in our
experiments for the task of declipping on SignalTrain.

###  6.4 Discussion

Figure 5: Spectrograms of a clipped guitar signal
($\gamma\\!=\\!$45\text{\,}\mathrm{dB}$$), its target signal, and the output
of each method. The respective time domain input and target signal are shown
at the top. While all methods significantly reduce the harmonics, it can be
observed that IRM and UMX smear the transients. A-SPADE relies on the
strongest bins in the clipped spectrum, which leads to tonal artifacts.

Qualitative Evaluation Although the abundance of evaluation metrics in the
literature has the potential to analyze the results in very detailed ways, it
does not always aid the judgement of which method is best given the individual
context. Often, different applications have different requirements concerning
the sound quality and can afford some types of distortions or artifacts to be
left in the signal. Therefore, we report some qualitative considerations that
need to be taken into account concerning our results, with the aim of finding
some descriptive patterns among them. Fig. 5 shows spectrograms of a hard-
clipped guitar signal and the outputs of all models under evaluation.

We found that the characteristics of the artifacts that each model produces
are consistent, independent from the task it is applied to. Nevertheless, for
the time domain-based models, the artifacts are less prominent in the de-
overdrive task. Here, Demucs often produces outputs that are virtually
indistinguishable from the original signal.

Generally, Demucs is the model that most often produces high quality results.
Especially for inputs with a low amount of distortion, it can reconstruct the
original sound without any perceptual artifact. Wave-U-Net behaves similarly,
although it often cannot reach the same quality.

UMX generally removes the distortion characteristics very well at the cost of
strong phasing artifacts: Despite the absence of input distortions, the
spectral features of the output are not necessarily consistent with the ones
of the target. This is most likely due to UMX re-using the phase of the
distorted input to reconstruct the signal in the time domain and the frame-
wise processing. Moreover, Fig. 5 highlights that the transients are smeared
by re-using the phase of the degraded signal.

CRAFx does not suffer from phasing artifacts, but occasionally leaves part of
the distortion features in the output. In some cases, the model fails to
reconstruct the onset of some notes, penalizing the listening experience.

Finally, A-SPADE is the model that exhibits the strongest and most frequent
artifacts, especially for strongly clipped signals. Although it considers
phase information by working in the complex frequency domain, it leads to non-
optimal solutions. Nevertheless, distortion features (like those left by
Demucs) or transient smearing (left by UMX) do not occur.

Influence of the Dry Signal We have observed a substantial difference in
performance between models that need to remove overdrive (CEG-OD dataset) and
models that need to remove hard-clipping (CEG-HC dataset). The superior
performance of the models trained and tested on CEG-OD can be mainly traced
back to the presence of the non-distorted signal in the overdrive output and
not to the soft-clipping character of the specific overdrive implementation.
We verified this hypothesis by training Wave-U-Net on hard-clipped data
superimposed with the clean signal (even when the amplitude of the clean
signal is considerably low). We obtained results similar to those in the de-
overdrive task (median results without superposition: $\textrm{SI-
SDR}=$7.2\text{\,}\mathrm{dB}$$, with superposition: $\textrm{SI-
SDR}=$34.4\text{\,}\mathrm{dB}$$). While the performance drop is present for
all metrics, it is less pronounced for UMX which seems not to utilize the
additional information in the signal.

Inference Speed The benefits of Demucs in the context of declipping go beyond
the quality of its outputs: using a neural approach also has advantages
regarding inference speed. Inference with A-SPADE is comparably slow (real-
time factor on CPU $\times\\!\mathrm{RT}\in[4.2,27.3]$ depending on
$\mathrm{SDR_{inp}}$ [13]), being an iterative approach that requires a
computation of the Fourier transform and its inverse at each iteration. Demucs
($\times\\!\mathrm{RT}=0.072$), Wave-U-Net ($\times\\!\mathrm{RT}=0.113$) and
UMX ($\times\\!\mathrm{RT}=0.026$) instead, allow for fast inference on the
CPU and even surpass real-time constraints independently on
$\mathrm{SDR_{inp}}$ without sacrificing the quality of the results.

Evaluation Metrics The results highlight how our evaluation metrics focus on
different aspects of the reconstructions: While SI-SDR measures differences
between two audio signals in the time domain, PEAQ focuses on the perceptual
quality without differentiating between degradation related to nonlinear
distortions and artifacts/quality. In contrast, $R_{\mathrm{nonlin}}$
specifically highlights nonlinear distortions that have not been removed.
Finally, FAD focuses primarily on the degradation that is observable in the
mel spectrum. Hence, each metric proves beneficial in measuring specific
aspects in the analysis of audio effect removal systems and can be advised for
further investigations.

##  7 CONCLUSION

We showed that recovering the clean signal from clipped or overdriven audio
signals can efficiently solved with neural networks designed for source
separation. We found that Demucs achieves high quality according to the chosen
evaluation metrics, especially when the distortion algorithm to be removed
blends the distorted sound with the original one. This outcome highlights the
potential of the proposed approach for other audio effects that mix dry and
wet signals (e.g., parallel compression, reverberation, delay, modulation
effects).

Moreover, we showed that Demucs, Wave-U-Net, and UMX outperform one state-of-
the-art declipping method on our test data. This outcome is promising,
considering that the dataset to train such a system is potentially much larger
than the size of the dataset we used. By discussing the results, we stressed
the usefulness of multiple evaluation metrics suitable to assess distortion
removal systems.

Future work should include gathering more clean electric guitar data and
generating a dataset using high-quality distortion emulations, which is
required to improve generalization on real-world data. Furthermore, the
knowledge from sparsity-based declipping algorithms could yield a valuable
prior for declipping through DNNs.

## References

  * [1] U. Zölzer, _DAFX: digital audio effects_ , 2nd ed. Chichester, West Sussex, England: Wiley, 2011.
  * [2] R. Stables, B. De Man, and J. D. Reiss, “Ten years of automatic mixing,” in _Proceedings of the 3rd Workshop on Intelligent Music Production, Salford, UK, 15 September 2017_ , 2017\.
  * [3] M. A. Martínez-Ramírez, D. Stoller, and D. Moffat, “A Deep Learning Approach to Intelligent Drum Mixing With the Wave-U-Net,” _Journal of the Audio Engineering Society_ , vol. 69, no. 3, pp. 142–151, Mar. 2021.
  * [4] C. J. Steinmetz, J. Pons, S. Pascual, and J. Serra, “Automatic Multitrack Mixing With A Differentiable Mixing Console Of Neural Audio Effects,” in _ICASSP 2021 - 2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. Toronto, ON, Canada: IEEE, Jun. 2021, pp. 71–75.
  * [5] Y. Mitsufuji, G. Fabbro, S. Uhlich, F.-R. Stöter, A. Défossez, M. Kim, W. Choi, C.-Y. Yu, and K.-W. Cheuk, “Music Demixing Challenge 2021,” _Frontiers in Signal Processing_ , vol. 1, Jan. 2022.
  * [6] W. Choi, M. Kim, M. A. Martínez-Ramírez, J. Chung, and S. Jung, “AMSS-Net: Audio Manipulation on User-Specified Sources with Textual Queries,” in _Proceedings of the 29th ACM International Conference on Multimedia_ , Apr. 2021, pp. 1775–1783.
  * [7] Z. Rafii, A. Liutkus, F.-R. Stöter, S. I. Mimilakis, and R. Bittner, “The MUSDB18 corpus for music separation,” Dec. 2017.
  * [8] T. Wilmering, D. Moffat, A. Milo, and M. B. Sandler, “A History of Audio Effects,” _Applied Sciences_ , vol. 10, no. 3, p. 791, Jan. 2020.
  * [9] M. A. Martínez-Ramírez, E. Benetos, and J. Reiss, “Deep Learning for Black-Box Modeling of Audio Effects,” _Applied Sciences_ , vol. 10, p. 638, Jan. 2020.
  * [10] C. Bagwell, et al., “Sound eXchange (SoX),” Feb. 2015.
  * [11] P. Záviška, P. Rajmic, A. Ozerov, and L. Rencker, “A Survey and an Extensive Evaluation of Popular Audio Declipping Methods,” _IEEE Journal of Selected Topics in Signal Processing_ , vol. 15, no. 1, pp. 5–24, Jan. 2021.
  * [12] A. Janssen, R. Veldhuis, and L. Vries, “Adaptive interpolation of discrete-time signals that can be modeled as autoregressive processes,” _IEEE Transactions on Acoustics, Speech, and Signal Processing_ , vol. 34, no. 2, pp. 317–330, Apr. 1986.
  * [13] C. Gaultier, S. Kitić, R. Gribonval, and N. Bertin, “Sparsity-Based Audio Declipping Methods: Selected Overview, New Algorithms, and Large-Scale Evaluation,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , vol. 29, pp. 1174–1187, 2021.
  * [14] H. B. Kashani, M. M. Goodarzi, A. Jodeiri, and S. G. Firooz, “Image to Image Translation based on Convolutional Neural Network Approach for Speech Declipping,” _4th Conference on Technology In Electrical and Computer Engineering (ETECH 2019)_ , 2019.
  * [15] O. Ronneberger, P. Fischer, and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” in _Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015_ , N. Navab, J. Hornegger, W. M. Wells, and A. F. Frangi, Eds. Cham: Springer International Publishing, 2015, pp. 234–241.
  * [16] W. Mack and E. A. P. Habets, “Declipping Speech Using Deep Filtering,” in _2019 IEEE Workshop on Applications of Signal Processing to Audio and Acoustics (WASPAA)_ , Oct. 2019, pp. 200–204.
  * [17] T. Tanaka, K. Yatabe, M. Yasuda, and Y. Oikawa, “APPLADE: Adjustable Plug-and-Play Audio Declipper Combining DNN with Sparse Optimization,” in _ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. Singapore, Singapore: IEEE, May 2022, pp. 1011–1015.
  * [18] M. A. Martínez-Ramírez, E. Benetos, and J. D. Reiss, “A general-purpose deep learning approach to model time-varying audio effects,” in _22nd International Conference on Digital Audio Effects (DAFx-19)_ , Jun. 2019.
  * [19] F.-R. Stöter, S. Uhlich, A. Liutkus, and Y. Mitsufuji, “Open-Unmix - A Reference Implementation for Music Source Separation,” _Journal of Open Source Software_ , vol. 4, no. 41, p. 1667, Sep. 2019.
  * [20] S. Uhlich, M. Porcu, F. Giron, M. Enenkl, T. Kemp, N. Takahashi, and Y. Mitsufuji, “Improving music source separation based on deep neural networks through data augmentation and network blending,” in _2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. New Orleans, LA: IEEE, Mar. 2017, pp. 261–265.
  * [21] S. Hochreiter and J. Schmidhuber, “Long Short-Term Memory,” _Neural Computation_ , vol. 9, no. 8, pp. 1735–1780, Nov. 1997.
  * [22] A. Narayanan and D. Wang, “Ideal ratio mask estimation using deep neural networks for robust speech recognition,” in _2013 IEEE International Conference on Acoustics, Speech and Signal Processing_ , 2013, pp. 7092–7096.
  * [23] D. Stoller, S. Ewert, and S. Dixon, “Wave-U-Net: A Multi-Scale Neural Network for End-to-End Audio Source Separation,” in _Proceedings of the 19th ISMIR Conference, Paris, France, September 23-27, 2018_ , Paris, Jun. 2018.
  * [24] A. Défossez, N. Usunier, L. Bottou, and F. Bach, “Music Source Separation in the Waveform Domain,” _arXiv:1911.13254 [cs, eess, stat]_ , Apr. 2021.
  * [25] S. Kitić, N. Bertin, and R. Gribonval, “Sparsity and Cosparsity for Audio Declipping: A Flexible Non-convex Approach,” in _Latent Variable Analysis and Signal Separation_ , E. Vincent, A. Yeredor, Z. Koldovský, and P. Tichavský, Eds. Cham: Springer International Publishing, 2015, pp. 243–250.
  * [26] M. Stein, J. Abeßer, C. Dittmar, and G. Schuller, “Automatic Detection of Audio Effects in Guitar and Bass Recordings,” in _AES 128th Convention_ , London, UK, May 2010.
  * [27] S. H. Hawley, B. Colburn, and S. I. Mimilakis, “SignalTrain: Profiling Audio Compressors with Deep Neural Networks,” in _AES 147th Convention_ , May 2019.
  * [28] D. P. Kingma and J. Ba, “Adam: A Method for Stochastic Optimization,” in _3rd International Conference on Learning Representations,{ICLR} 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_, 2015.
  * [29] E. Vincent, R. Gribonval, and C. Fevotte, “Performance measurement in blind audio source separation,” _IEEE Transactions on Audio, Speech and Language Processing_ , vol. 14, no. 4, pp. 1462–1469, Jul. 2006.
  * [30] J. L. Roux, S. Wisdom, H. Erdogan, and J. R. Hershey, “SDR – Half-baked or Well Done?” in _ICASSP 2019 - 2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. Brighton, United Kingdom: IEEE, May 2019, pp. 626–630.
  * [31] E. Cano, D. FitzGerald, and K. Brandenburg, “Evaluation of quality of sound source separation algorithms: Human perception vs quantitative metrics,” in _2016 24th European Signal Processing Conference (EUSIPCO)_. Budapest, Hungary: IEEE, Aug. 2016, pp. 1758–1762.
  * [32] V. Emiya, E. Vincent, N. Harlander, and V. Hohmann, “Subjective and Objective Quality Assessment of Audio Source Separation,” _IEEE Transactions on Audio, Speech, and Language Processing_ , vol. 19, no. 7, pp. 2046–2057, Sep. 2011.
  * [33] T. Thiede, W. C. Treurniet, R. Bitto, C. Schmidmer, T. Sporer, J. G. Beerends, and C. Colomes, “PEAQ - The ITU Standard for Objective Measurement of Perceived Audio Quality,” _Journal of the Audio Engineering Society_ , vol. 48, no. 1/2, pp. 3 EP – 29, Feb. 2000.
  * [34] S. Lattner and J. Nistal, “Stochastic Restoration of Heavily Compressed Musical Audio Using Generative Adversarial Networks,” _Electronics_ , vol. 10, no. 11, p. 1349, Jun. 2021.
  * [35] J. You, U. Reiter, M. M. Hannuksela, M. Gabbouj, and A. Perkis, “Perceptual-based quality assessment for audio–visual services: A survey,” _Signal Processing: Image Communication_ , vol. 25, no. 7, pp. 482–501, Aug. 2010.
  * [36] C.-T. Tan, B. C. J. Moore, N. Zacharov, and V.-V. Mattila, “Predicting the Perceived Quality of Nonlinearly Distorted Music and Speech Signals,” _Journal of the Audio Engineering Society_ , vol. 52, no. 7/8, pp. 699–711, Jul. 2004.
  * [37] K. Kilgour, M. Zuluaga, D. Roblek, and M. Sharifi, “Fréchet Audio Distance: A Reference-Free Metric for Evaluating Music Enhancement Algorithms,” in _Interspeech 2019_. ISCA, Sep. 2019, pp. 2350–2354.
  * [38] S. Hershey, S. Chaudhuri, D. P. W. Ellis, J. F. Gemmeke, A. Jansen, R. C. Moore, M. Plakal, D. Platt, R. A. Saurous, B. Seybold, M. Slaney, R. J. Weiss, and K. Wilson, “CNN Architectures for Large-Scale Audio Classification,” _International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , Jan. 2017.
  * [39] D. Dowson and B. Landau, “The Fréchet distance between multivariate normal distributions,” _Journal of multivariate analysis_ , vol. 12, no. 3, pp. 450–455, 1982.
  * [40] W. Choi, M. Kim, J. Chung, D. Lee, and S. Jung, “Investigating U-Nets with various Intermediate Blocks for Spectrogram-based Singing Voice Separation,” in _Proceedings of the 21th International Society for Music Information Retrieval Conference_ , Oct. 2020.
  * [41] A. Défossez, “Hybrid Spectrogram and Waveform Source Separation,” in _Proceedings of the ISMIR 2021 Workshop on Music Source Separation_ , Nov. 2021.


(a) shows the similarity between the representations learned by different layers of the encoder, transition and critic networks of a CHMM whose critic maximises the reward (with softmax action selection). (b) shows the similarity between the representations learned by different layers of the encoder, transition and critic networks of a CHMM whose critic minimises the EFE (with softmax action selection). (c) shows the similarity between the representations learned by two CHMMs, one whose critic optimises EFE and the other that optimises reward (both with softmax action selection).

in which all our simulations will be

(a) shows one latent dimension of the variance layer of the transition network for the CHMM maximising the reward.
(b) shows one latent dimension of the variance layer of the transition network for the CHMM minimising the EFE.
Both figures are representative of the distributions of variance activations in the two models.
Note, only the action down has low variance for the CHMM minimising the EFE.
This suggests that the CHMM minimising the EFE always pick the action down, and does not gather enough data for the other actions.

§.§.§ Degenerate behaviour with the expected free energy?

Up to now, we saw that the CHMM minimising expected free energy (EFE) was not able to solve the task. Also, we discovered that the transition network is uncertain for the actions: up, left, and right, which suggests that the CHMM minimising EFE always takes action down. Figure <ref> corroborates this story. Indeed, Figure <ref> shows that the agent minimising EFE almost exclusively picked action down, and Figure <ref> shows that the entropy of the prior over actions very quickly converges to zero.

In contrast, the CHMM maximising reward, keeps on selecting the actions right and left, which enables it to drag the shape towards the appropriate corner (see Figure <ref>). Also, as shown in Figure <ref>, the entropy of the prior over actions remained a lot higher than zero. Note, the only difference between the CHMM minimising EFE and the one maximising reward is the information gain, which is defined as the KL divergence between the output of the transition and encoder networks. Since the EFE is minimised, the output of those two networks need to be as close as possible to each other.

(a) shows the action taken for each planning iteration when the CHMM is minimising expected free energy.
(b) shows the entropy of the prior over actions when the CHMM is minimising the EFE.
(c) shows the action taken for each planning iteration when the CHMM is maximising reward.
(d) shows the entropy of the prior over actions when the CHMM is maximising reward.

This suggests that the CHMM minimising EFE is picking a single action (down), and becomes an expert at predicting the future when selecting this action. This effectively makes the KL divergence between the output of the transition and encoder networks small. Additionally, when selecting the action down, the average reward is zero, because (in the dSprites dataset) there are as many shapes on the left of the image as on the right, and when crossing the bottom line, the agent receives a reward which is linearly increasing (or decreasing) as a corner is approached and is zero at the center of the image. For all the other actions, the expected reward will be negative because after 50 action-perception cycles without crossing the bottom line, the trial is interrupted and the agent receive a reward of -1. Thus, if the CHMM has to stick to a single action to keep the KL divergence  small, then the best action it can choose is down, i.e., action down has the highest expected reward.

Also, Figure <ref> shows the impact of adding X% of the information gain into the objective function, i.e., the agent starts by only maximising reward (c.f. Equation <ref>), and after 200K training iterations minimises reward plus X% of the information gain. One can see that adding even 1% of the information gain already dramatically decreases the amount of reward gathered.

To conclude, the same information gain that is intended to give an EFE minimising agent its exploration behaviour, also prevents the agent from solving the dSprites environment. This is because the agent is reduced to picking a single action, leading to a suboptimal policy.

[square/.style=regular polygon,regular polygon sides=4, scale=1]
at (-2, 0) ;

[black] at (4.675, -1) CHMM[0%]:;
[fill=DarkBlueCHMM] (6, -0.8) rectangle (7, -1.2);

[black] at (4.675, -0.5) CHMM[1%]:;
[fill=GrayCHMM] (6, -0.7) rectangle (7, -0.3);

[black] at (4.675, 0) CHMM[5%]:;
[fill=OrangeCHMM] (6, -0.2) rectangle (7, 0.2);

[black] at (4.585, 0.5) CHMM[15%]:;
[fill=GreenCHMM] (6, 0.3) rectangle (7, 0.7);

[black] at (4.585, 1) CHMM[25%]:;
[fill=PinkCHMM] (6, 0.8) rectangle (7, 1.2);

[black] at (4.585, 1.5) CHMM[50%]:;
[fill=LightBlueCHMM] (6, 1.3) rectangle (7, 1.7);

[black] at (-2.8, -3.5) Training iterations;
[black, rotate=90] at (-7.3, 0) Total Reward;

This figure illustrates the total reward aggregated by CHMM agents during the 500K iterations of training. All the agents start by only maximising reward, and after 200K training iterations, X% of the information gain is added to the objective function. Note, even adding 1% of the information gain is enough to drastically reduce the total reward aggregated by the agent. The differences in trajectories before 200K are arbitrary, arising from differences in random initializations.

§.§ DAI agent

In this section, we report the results obtained by the DAI agent, when using different action selection strategies and different definitions of the expected free energy. First, most of the fifteen DAI agents crashed because of numerical instability, i.e., the VFE suddenly became “Not a Number". The only DAI agent that survived (i.e., did not crash) was maximising rewards while performing softmax sampling for action selection. Figure <ref> shows that the DAI agent successfully minimises its variational free energy, but as shown in Figure <ref>, the DAI agent does not solve the task and performs as well as a random agent. Finally, Figure <ref> shows sequences of images produced by the DAI agent after 500K training iterations. Note, while the agent does not solve that task, it understands the dynamics of the environment pretty well. However, the agent struggles with images representing hearts.

By comparing Figures <ref>, <ref> and <ref>, we see that the DAI agent with softmax sampling has a better reconstruction than the CHMM agent with the $\mathring{\epsilon}$-greedy algorithm (which is presented in Figure <ref>). In contrast, the DAI agent does not reconstruct the sequences of images as well as the HMM agent performing random actions (which is presented in Figure <ref>).

[square/.style=regular polygon,regular polygon sides=4, scale=0.75]
at (-0.5, 0) ;
[gray!60!black] at (8, 0) Ground Truth (GT):;
[gray!60!black] at (8.1, 0.5) Reconstruction (R):;
[fill=Green] (10.5, 0.2) rectangle (11.5, -0.2);
[fill=Red] (10.5, 0.3) rectangle (11.5, 0.7);

This figure illustrates the sequences of reconstructed images generated by the DAI after 500K training iterations. The columns alternate between the ground truth images and the reconstructed images. Time passes vertically (from top to bottom), and within each column, the same action is executed repeatedly.

(a) shows the similarity between the representations learned by a DQN and a CHMM maximising the reward and using softmax action selection. (b) shows the similarity between the representations learned by a DQN and a DAI maximising the reward and using softmax action selection. (c) shows the similarity between the representations learned by a CHMM and a DAI. Both maximise the reward and use softmax action selection.

As previously mentioned, the only DAI that did not crash during training used softmax action selection and had a critic maximising reward. We can see in Figures <ref> and <ref> that the representational similarity between this DAI and DQN is very close to the representational similarity between a DQN and a CHMM using the same action selection and maximising reward. This is further confirmed by a comparison between the CHMM and the DAI model in Figure <ref>. Interestingly, we can see that the policy and critic network learn similar representations, indicating that the policy network is learning correctly. However, we previously inferred that the softmax action selection may be suboptimal and this seems to hold true for the DAI as well, given that it is unable to solve the task.

[square/.style=regular polygon,regular polygon sides=4, scale=0.75]
at (0, 0) ;
[gray!60!black, rotate=90] at (-7.7, 0) Total rewards;
[gray!60!black] at (0, -4.8) Training Iterations;
[RedDQN, anchor=west] at (7.2, 4) DQN;
[BlueCHMM, anchor=west] at (7.2, -2.6) $DAI[G^4]$;

This figure illustrates the total amount of reward gathered by a DAI agent during the 500K iterations of training. This agent was maximising rewards while sampling actions from a softmax function of the policy network output. Put simply, the DAI agent does not solve the task and performs at the level of a random agent.

[square/.style=regular polygon,regular polygon sides=4, scale=0.75]
at (0, 0) ;
[gray!60!black, rotate=90] at (-8, 0) VFE;
[gray!60!black] at (0, -6.7) Training Iterations;

This figure illustrates the variational free energy of the DAI agent during the 500K iterations of training. This agent was maximising reward while sampling actions from a softmax function of the policy network output. The agent was able to minimise its variational free energy.

§ DISCUSSION OF EPISTEMIC VALUE

In this section, we discuss the decomposition of the expected free energy into epistemic and extrinsic value. More precisely, intrigued by the poor results of deep active inference agents (e.g., CHMM and DAI), we seek to understand the damaging impact of the epistemic value on performance. Assuming the following definition for the epistemic value:

\begin{align*}
\bm{EV} &= \mathbb{E}_{P(o_\tau|s_\tau)Q(s_\tau | \pi)}\big[\ln P(s_\tau|o_\tau) - \ln Q(s_\tau | \pi) \big],
\end{align*}

we set up two experiments (c.f., Appendix C for details) in which all the disributions are stored in tables. In other words, the following categorical distributions are represented using matrices: $P(o_\tau|s_\tau)$, $Q(s_\tau | \pi)$, and $P(s_\tau|o_\tau)$.

In the first experiement, the prior over states $P(s_\tau)$ is fixed, and the likelihood $P(o_\tau|s_\tau)$ is becoming more and more uniform. While this is happening, the true posterior $P(s_\tau|o_\tau)$ becomes more similar to the approximate posterior $Q(s_\tau | \pi)$, and the epistemic value decreases; see left panel of Figure <ref>. This is the expected behaviour, and in this setting, the epistemic value encourages exploration.

In the second experiment, the likelihood $P(o_\tau|s_\tau)$ is fixed and has rather high entropy, while the prior over states $P(s_\tau)$ is shifted in one direction on the state axis. As a result of this, the true posterior $P(s_\tau|o_\tau)$ becomes more different to the approximate posterior $Q(s_\tau | \pi)$, but this causes the epistemic value to decrease; see Figure <ref>, right panel. This is a degenerate behaviour, as in this setting, the epistemic value discourages exploration.

To sum up, the expected free energy decomposition into epistemic and extrinsic value — as presented in Equation (10) of Parr and Friston, 2019 — seems to exhibit two very different behaviours depending on how the distributions are defined, c.f., Figure <ref>. This is particularly important in the deep active inference literature, which builds on this equation. For example, the graphs presented in Section <ref> indicate that the CHMM agent minimising expected free energy is focusing almost exclusively on the action “down", i.e., it is not exploring, which leads to poor performance.

[square/.style=regular polygon,regular polygon sides=4]
at (0,0) ;
[rotate=90] at (-4.3,0) Epistemic Value;
at (0.2,-3.5) Similarity of posteriors: $P(s_\tau|o_\tau)$ and $Q(s_\tau|\pi)$;

[square/.style=regular polygon,regular polygon sides=4]
at (0,0) ;
[rotate=90] at (-4.3,0) Epistemic Value;
at (0.05,-3.5) Similarity of posteriors: $P(s_\tau|o_\tau)$ and $Q(s_\tau|\pi)$;

The left-most figure shows the result of the first experiment where the likelihood was becoming more and more uniform. In this setting, the epistemic value encourages exploration, as maximising the epistemic value makes the true posterior $P(s_\tau | o_\tau)$ as different as possible to the approximate posterior $Q(s_\tau | \pi)$, leftward on x-axis. However, the right-most figure illustrates the result of the second experiment where the prior over states was changing. In this case, the epistemic value does not promote exploration, as maximising the epistemic value makes the true posterior $P(s_\tau | o_\tau)$ as similar as possible to the approximate posterior $Q(s_\tau | \pi)$, rightward on x-axis. Indeed, in this case, maximising the epistemic value causes information to be lost.

§ CONCLUSION

In this paper, we challenged the common assumption that deep active inference is a solved problem, by highlighting the challenges that need to be resolved for the field to move forward. We reviewed eight approaches implementing deep active inference: (a) $DAI_{MC}$ by Fountas et al., 2020, (b) $DAI_{VPG}$ by Millidge, 2020, (c) $DAI_{RHI}$ by Rood et al., 2020, (d) $DAI_{HR}$ by Sancaktar et al., 2020, (e) $DAI_{FA}$ by Ueltzhöffer, 2018, (f) $DAI_{POMDP}$ by van der Himst and Lanillos, 2020, (g) $DAI_{SSM}$ by Çatal et al., 2020, and (h) the approach by Schneider et al., 2022 for which the code was not available online.

Overall, those approaches brought interesting ideas such as: using deep neural networks to predict the parameters of the distributions of interest, using Monte-Carlo tree search for planning in active inference, and using a bootstrap estimate of the expected free energy to train a critic network. Yet, we struggled to replicate some of the results claimed, e.g. training $DAI_{MC}$ on the animal AI environment, and we were unable to access code in some instances. Ideally, future research should draw inspiration from the open science framework, by making the code that produced the claimed results open source. Also, the definition of the expected free energy varied between papers. This suggests that additional research is required to clarify the definition and justification of the expected free energy both in tabular and deep active inference (c.f., Section <ref> and Appendix C). To sum up, recent research on deep active inference [Fountas et al., 2020, Millidge, 2020, Rood et al., 2020, Sancaktar et al., 2020, Ueltzhöffer, 2018, van der Himst and Lanillos, 2020, Çatal et al., 2020] has made an important first step towards a complete deep active inference agent, but a number of details still need to be honed, e.g., the definition and derivation of the expected free energy, and reproducibilty of research. For more details, the reader is referred to Section <ref>.

After reviewing existing research, we tried to progressively implement a deep active inference agent. First, we produced a variational auto-encoder agent that takes random actions. This agent was able to learn a useful (latent) representation of the task. However, since the agent takes random actions, it was unable to solve the task. Then, we added the transition network to create the hidden Markov model (HMM) agent, which also takes random actions. The agent was able to learn a good represenation of the task and of its dynamics, but was not able to solve the task.

Next, we tried to incorporate a critic network into the approach leading to the critical hidden Markov model (CHMM). In this context, we experimented with several possible implementations of the expected free energy. We also tried to remove the information gain and simply predict the reward. Additionally, we implemented three types of action selection strategies, namely: best action according to the expected free energy, softmax sampling, and epsilon-greedy.

When the epsilon-greedy algorithm was used, only the agent maximising reward was able to solve the task.
However, the agent requires more training iterations than a simple deep Q-network to learn the right behaviour.
This may be explained by the fact that the CHMM not only has to learn to solve the task, but also learn the dynamics of the environment.
When softmax sampling was used, all the agents failed to solve the task. One of them perfoming even worse than an agent selecting random actions.
Lastly, when selecting the best action, only the reward maximising agent was able to solve the task.
Importantly, according to our experiments, the agent using the epsilon-greedy algorithm received the highest amount of cumulated rewards and learned to solve the task the fastest.
Additionally, the reward maximising agents properly solve the task, but the quality of their latent representation is not as good as that of an HMM agent.
This may be due to the fact that when performing reward maximising actions, the data available to learn the model of the environment lacks diversity, i.e., not enough exploration.

Next, we tried to incorporate a policy network into the approach leading to a complete deep active inference agent (DAI). As for the CHMM agent, we experiemented with several possible implementations of the expected free energy, tried to remove the information gain and simply predict the reward, and implemented three types of action selection strategies, namely: best action according to the expected free energy, softmax sampling, and epsilon-greedy algorithm. When the epsilon-greedy algorithm was used or the best action was selected, all agents failed to solve the task. When using softmax sampling, most of the agents were numerically unstable and crashed, and the remaining agents failed to solve the task.

Finally, we compared the similarity of the representation learned by the layers of various models (e.g., deep Q-network, CHMM, DAI, etc...) using centered kernel alignment. This reveals that the DQN learns general features in its first few layers, and very specialised features in its last two layers. The VAE learns similar features to the DQN in the first layers, but differs from the DQN in the last two layers, reflecting the difference in learning objectives. Similarly, the HMM learns similar features to the DQN in the first layers, but differs from the representation learned by the DQN in the last two layers. Also, the mean and variance representations learned by the HMM are different from their VAE counterparts, which suggests that the transition network influences the latent representation of the model.

Additionally, when using the $\mathring{\epsilon}$-greedy algorithm for action selection, the representations learned by the CHMM maximising reward is closer to the DQN than the CHMM minimising expected free energy is to the DQN. Importantly, the critic network of the reward maximising CHMM retains more information from the encoder than the CHMM minimising expected free energy. When the best action (according to the critic network) is selected, the CHMM maximising reward and the CHMM minimising expected free energy learn very similar representations except for the variance layers of the transition and encoder network. While performing further inspection of those (variance) layers, we found that the transition network of the reward maximising CHMM is a lot more certain than the transition network of the CHMM minimising expected free energy. More precisely, the CHMM minimising expected free energy is only confident about the world transition when performing action down. This suggests that the CHMM minimising expected free energy always picks the action down, and does not gather enough data for the other actions. Visualising the distribution of actions selected as training progresses corroborates this story by showing that the agent minimising EFE almost exclusively picks action down. In contrast, the CHMM maximising reward, keeps on selecting the actions left and right, which enables it to successfully solve the task. The only difference between those two CHMMs is the epistemic value, which aims to make the outputs of the transition and encoder network as close as possible. Thus, the CHMM minimising expected free energy is picking a single action (down), and becomes an expert at predicting the future when selecting this action. This effectively makes the KL divergence between the output of the transition and encoder networks small. Additionally, when selecting the action down, the average reward is zero, while for all the other actions, the expected reward will be
negative. Therefore, if the CHMM has to stick to a single action to keep the KL divergence small, then the action down
is the most rewarding.

The same observations about the variance layers also applies to CHMMs using softmax sampling for action selection. While this may explain why the model optimising the expected free energy does not solve the task, it does not explain why the model maximising the reward cannot solve the task, and we can hypothesise that those results may be due to the softmax action selection. More precisely, if the values predicted by the critic network are very close to each other, then an agent using softmax sampling may perform random actions. Note, increasing the gain parameter may help the agent to differentiate between values close to each other.

In addition, the representational similarity between the DAI (maximising reward using softmax sampling) and DQN is very close to the representational similarity between a DQN with a CHMM (maximising reward using softmax sampling). Also, the DAI's policy and critic network learn similar representations, which indicates that the policy network is learning correctly. Thus, the fact that the DAI (maximising reward using softmax sampling) fails in the dSprites environment is likely due to the softmax action selection and not to the representation learned by the model.

Lastly, our investigations of the expected free energy often used in deep active inference [Fountas et al., 2020], suggests that degenerate behaviour can arise from it, in certain situations. This could explain why adding epistemic value to our planning objective seems to have such a damaging impact on the agent's capacity to explore its environment and gain information. The use of this definition of the expected free energy in the deep learning context may be at the core of our difficulty getting deep active inference to work and may also explain some of the presentational uncertainties (e.g., additions of minus signs) found in the deep active inference literature.

To conclude, the field of deep active inference has benefited from a large variety of ideas from the reinforcement and deep learning literature. In the future, it would be valuable to provide an approach that satisfies the following five desirata: (i) the approach is complete, i.e., it is composed of an encoder, a decoder, a transition network, a policy network and (optionally) a critic network, (ii) the mathematics underlying the approach is errorless and consistent with the free energy principle, (iii) the implementation is consistent with the mathematics, (iv) the code is publicly available so that the correctness of the implementation can be verified and the results reproduced, and (v) the approach is able to solve tasks with a large input space, e.g. image-based tasks. We believe that such an approach will benefit the field of deep active inference by providing a strong and reproducible baseline against which future research could benchmark.

We thank Karl J. Friston, Thomas Parr, Lancelot Da Costa, and Zafeirios Fountas for useful discussions and feedback surronding this paper, as well as for pointing us towards important resources such as their new book and several papers.

[Bellman, 1952]
Richard Bellman.
On the theory of dynamic programming.
Proc Natl Acad Sci U S A, 380 (8):0 716–719,
August 1952.

[Bonheme and Grzes, 2022]
Lisa Bonheme and Marek Grzes.
How do Variational Autoencoders Learn? Insights from
Representational Similarity.
arXiv e-prints, art. arXiv:2205.08399, May 2022.

[Browne et al., 2012]
C. B. Browne, E. Powley, D. Whitehouse, S. M. Lucas, P. I. Cowling,
P. Rohlfshagen, S. Tavener, D. Perez, S. Samothrakis, and
S. Colton.
A survey of monte carlo tree search methods.
IEEE Transactions on Computational Intelligence and AI in
Games, 40 (1):0 1–43, 2012.

[Çatal et al., 2020]
Ozan Çatal, Samuel Wauthier, Cedric De Boom, Tim Verbelen, and Bart
Learning generative state space models for active inference.
Frontiers in Computational Neuroscience, 14:0 574372,

[Champion et al., 2021]
Théophile Champion, Marek Grześ, and Howard Bowman.
Realizing Active Inference in Variational Message Passing: The
Outcome-Blind Certainty Seeker.
Neural Computation, 330 (10):0 2762–2826, 09
ISSN 0899-7667.
URL <https://doi.org/10.1162/neco_a_01422>.

[Champion et al., 2022]
Théophile Champion, Howard Bowman, and Marek Grześ.
Branching time active inference: Empirical study and complexity class
Neural Networks, 152:0 450–466, 2022a.
ISSN 0893-6080.

[Champion et al., 2022]
Théophile Champion, Lancelot Da Costa, Howard Bowman, and Marek Grześ.
Branching time active inference: The theory and its generality.
Neural Networks, 151:0 295–316, 2022b.
ISSN 0893-6080.

[Champion et al., 2022]
Théophile Champion, Marek Grześ, and Howard Bowman.
Multi-modal and multi-factor branching time active inference,
URL <https://arxiv.org/abs/2206.12503>.

[Champion et al., 2022]
Théophile Champion, Marek Grześ, and Howard Bowman.
Branching Time Active Inference with Bayesian Filtering.
Neural Computation, 340 (10):0 2132–2144, 09
ISSN 0899-7667.
URL <https://doi.org/10.1162/neco_a_01529>.

[Cortes et al., 2012]
Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh.
Algorithms for Learning Kernels Based on Centered Alignment.
J. Mach. Learn. Res., 130 (1):0 795–828,
March 2012.
ISSN 1532-4435.

[Costa et al., 2020]
Lancelot Da Costa, Thomas Parr, Noor Sajid, Sebastijan Veselic, Victorita
Neacsu, and Karl Friston.
Active inference on discrete state-spaces: a synthesis,

[Costa et al., 2020]
Lancelot Da Costa, Noor Sajid, Thomas Parr, Karl Friston, and Ryan Smith.
The relationship between dynamic programming and active inference:
the discrete, finite-horizon case, 2020b.

[Cristianini et al., 2002]
Nello Cristianini, John Shawe-Taylor, André Elisseeff, and Jaz S Kandola.
On Kernel-Target Alignment.
In T G Dietterich, S Becker, and Z Ghahramani, editors,
Advances in Neural Information Processing Systems, volume 14, pages
367–373. MIT Press, Vancouver, Canada, 2002.

[Cullen et al., 2018]
Maell Cullen, Ben Davey, Karl J. Friston, and Rosalyn J. Moran.
Active inference in openai gym: A paradigm for computational
investigations into psychiatric illness.
Biological Psychiatry: Cognitive Neuroscience and
Neuroimaging, 30 (9):0 809 – 818, 2018.
ISSN 2451-9022.
Computational Methods and Modeling in Psychiatry.

[Doersch, 2016]
Carl Doersch.
Tutorial on variational autoencoders, 2016.

[FitzGerald et al., 2015]
Thomas H. B. FitzGerald, Raymond J. Dolan, and Karl Friston.
Dopamine, reward learning, and active inference.
Frontiers in Computational Neuroscience, 9:0 136,
ISSN 1662-5188.

[Fountas et al., 2020]
Zafeirios Fountas, Noor Sajid, Pedro A. M. Mediano, and Karl Friston.
Deep active inference agents using Monte-Carlo methods, 2020.

[Friston et al., 2016]
Karl Friston, Thomas FitzGerald, Francesco Rigoli, Philipp Schwartenbeck,
John O Doherty, and Giovanni Pezzulo.
Active inference and learning.
Neuroscience & Biobehavioral Reviews, 68:0 862 –
879, 2016.
ISSN 0149-7634.

[Gretton et al., 2005]
Arthur Gretton, Olivier Bousquet, Alex Smola, and Bernhard Schölkopf.
Measuring Statistical Dependence with Hilbert-Schmidt Norms.
In Sanjay Jain, Hans Ulrich Simon, and Etsuji Tomita, editors,
Algorithmic Learning Theory, pages 63–77. Springer Berlin Heidelberg,
ISBN 978-3-540-31696-1.

[Higgins et al., 2017]
Irina Higgins, Loïc Matthey, Arka Pal, Christopher Burgess, Xavier
Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.
beta-VAE: Learning basic visual concepts with a constrained
variational framework.
In 5th International Conference on Learning Representations,
ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track
Proceedings. OpenReview.net, 2017.
URL <https://openreview.net/forum?id=Sy2fzU9gl>.

[Itti and Baldi, 2009]
Laurent Itti and Pierre Baldi.
Bayesian surprise attracts human attention.
Vision Research, 490 (10):0 1295 – 1306,
ISSN 0042-6989.
Visual Attention: Psychophysics, electrophysiology and neuroimaging.

[Kingma and Welling, 2014]
Diederik P. Kingma and Max Welling.
Auto-Encoding Variational Bayes.
In International Conference on Learning Representations,
volume 2, Banff, Canada, 2014.
URL <http://arxiv.org/abs/1312.6114>.

[Koller and Friedman, 2009]
Daphne Koller and Nir Friedman.
Probabilistic graphical models: principles and techniques.
MIT press, 2009.

[Kornblith et al., 2019]
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton.
Similarity of Neural Network Representations Revisited.
In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors,
Proceedings of the 36th International Conference on Machine Learning,
volume 97 of Proceedings of Machine Learning Research, pages
3519–3529, Long Beach, USA, Jun 2019. PMLR.
URL <http://proceedings.mlr.press/v97/kornblith19a.html>.

[Lample and Chaplot, 2016]
Guillaume Lample and Devendra Singh Chaplot.
Playing fps games with deep reinforcement learning, 2016.

[Lanillos et al., 2020]
Pablo Lanillos, Gordon Cheng, et al.
Robot self/other distinction: active inference meets neural networks
learning in a mirror.
arXiv preprint arXiv:2004.05473, 2020.

[Maheswaranathan et al., 2019]
Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David
Universality and individuality in neural dynamics across large
populations of recurrent networks.
In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural
Information Processing Systems, volume 32. Curran Associates, Inc., 2019.

[Matthey et al., 2017]
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner.
dsprites: Disentanglement testing sprites dataset.
https://github.com/deepmind/dsprites-dataset/, 2017.

[Millidge, 2019]
Beren Millidge.
Combining active inference and hierarchical predictive coding: A
tutorial introduction and case study., 2019.
URL <https://doi.org/10.31234/osf.io/kf6wc>.

[Millidge, 2020]
Beren Millidge.
Deep active inference as variational policy gradients.
Journal of Mathematical Psychology, 96:0 102348, 2020.
ISSN 0022-2496.

[Mnih et al., 2013]
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis
Antonoglou, Daan Wierstra, and Martin Riedmiller.
Playing atari with deep reinforcement learning, 2013.

[Oliver et al., 2019]
Guillermo Oliver, Pablo Lanillos, and Gordon Cheng.
Active inference body perception and action for humanoid robots.
arXiv preprint arXiv:1906.03022, 2019.

[Parr and Friston, 2019]
Thomas Parr and Karl J Friston.
Generalised free energy and active inference.
Biological cybernetics, 1130 (5-6):0 495–513,

[Pezzato et al., 2020]
Corrado Pezzato, Carlos Hernandez, and Martijn Wisse.
Active inference and behavior trees for reactive action planning and
execution in robotics, 2020.

[Rezende et al., 2014]
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra.
Stochastic Backpropagation and Approximate Inference in Deep
Generative Models.
In Eric P Xing and Tony Jebara, editors, Proceedings of the
31st International Conference on Machine Learning, volume 32 of
Proceedings of Machine Learning Research, pages 1278–1286, Bejing,
China, 2014. PMLR.
URL <http://proceedings.mlr.press/v32/rezende14.html>.

[Robert and Escoufier, 1976]
P. Robert and Y. Escoufier.
A Unifying Tool for Linear Multivariate Statistical Methods: The RV-
Journal of the Royal Statistical Society. Series C (Applied
Statistics), 250 (3):0 257–265, 1976.
ISSN 00359254, 14679876.
URL <http://www.jstor.org/stable/2347233>.

[Rood et al., 2020]
Thomas Rood, Marcel van Gerven, and Pablo Lanillos.
A deep active inference model of the rubber-hand illusion.
In Tim Verbelen, Pablo Lanillos, Christopher L. Buckley, and Cedric
De Boom, editors, Active Inference, pages 84–91, Cham, 2020. Springer
International Publishing.
ISBN 978-3-030-64919-7.

[Sancaktar et al., 2020]
Cansu Sancaktar, Marcel van Gerven, and Pablo Lanillos.
End-to-end pixel-based deep active inference for body perception and
action, 2020.

[Schneider et al., 2022]
Tim Schneider, Boris Belousov, Hany Abdulsamad, and Jan Peters.
Active inference for robotic manipulation.
arXiv preprint arXiv:2206.10313, 2022.

[Schwartenbeck et al., 2018]
Philipp Schwartenbeck, Johannes Passecker, Tobias U Hauser, Thomas H B
FitzGerald, Martin Kronbichler, and Karl Friston.
Computational mechanisms of curiosity and goal-directed exploration.
bioRxiv, 2018.
URL <https://www.biorxiv.org/content/early/2018/09/07/411272>.

[Silver et al., 2016]
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George
van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas
Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal
Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray
Kavukcuoglu, Thore Graepel, and Demis Hassabis.
Mastering the game of go with deep neural networks and tree search.
Nature, 5290 (7587):0 484–489, 2016.
URL <https://doi.org/10.1038/nature16961>.

[Sutton et al., 1998]
Richard S Sutton, Andrew G Barto, et al.
Introduction to reinforcement learning.
MIT press Cambridge, 1998.

[Ueltzhöffer, 2018]
Kai Ueltzhöffer.
Deep active inference.
Biol. Cybern., 1120 (6):0 547–573, dec 2018.
ISSN 0340-1200.
URL <https://doi.org/10.1007/s00422-018-0785-7>.

[van der Himst and Lanillos, 2020]
Otto van der Himst and Pablo Lanillos.
Deep active inference for partially observable mdps.
CoRR, abs/2009.03622, 2020.
URL <https://arxiv.org/abs/2009.03622>.

[van Hasselt et al., 2015]
Hado van Hasselt, Arthur Guez, and David Silver.
Deep reinforcement learning with double Q-learning, 2015.

left=0.8cm, bottom=2cm, top=1cm

§ APPENDIX A: NOTATION

Symbol     Meaning

$s_\tau$, $o_\tau$, $r_\tau$, $a_\tau$     State, observation, reward and action at time step $\tau$, respectively.

3*$o_\tau^r$, $\hat{s}_\tau^r$, $\mathring{o}^r_\tau$     An observation in which a reward has been encoded as explained in Figure <ref>, the hidden

state sampled from the encoder when feeding $o_\tau^r$ as input, and the observation reconstructed

by the decoder from the output of the transition network, respectively.

2*$\hat{s}_\tau$, $\hat{o}_\tau$     A state sampled from the encoder at time step $\tau$ when $o_\tau$ is provided as input, and the image

reconstructed by the decoder at time $\tau$ when using $\hat{s}_\tau$ as input, respectively.

$\mathring{o}_\tau$     The observations at time step $\tau$ predicted by the transition network.

$s_{i:j}$, $o_{i:j}$, $a_{i:j}$     Respectively, the set of states, observations, and actions between time step $i$ and $j$ (included).

2*$\pi$, $\pi_\tau$, $\pi'$     A policy, i.e. a sequence of actions, the action prescribed by the policy at time step $\tau$, and

another policy whose size is smaller or equal than the size of $\pi$, respectively.

$\mathcal{A}$, $\Pi$     The set of possible actions, and the set of possible policies, respectively.

$\mathbb{A}_s$     The set of all anscestors of a node $s$.

$\#A, \#C, \#S, \#O$     The number of actions, channels, states and observations, respectively.

$\mathcal{Q}_{\theta_a}$, $\hat{\mathcal{Q}}_{\hat{\theta}_a}$     The Q-network parameterised by $\theta_a$, and the target network parameterised by $\hat{\theta}_a$

$\mathcal{G}_{\theta_a}$, $\hat{\mathcal{G}}_{\hat{\theta}_a}$     The critic network parameterised by $\theta_a$, and the target network parameterised by $\hat{\theta}_a$.

$\mathcal{E}_{\phi_s}$, $\mathcal{D}_{\theta_o}$     The encoder network parameterised by $\phi_s$, and the decoder network parameterised by $\theta_o$.

$\mathcal{P}_{\phi_a}$     The policy network parameterised by $\phi_a$.

$\mathcal{T}_{\theta_s}$, $\mathcal{T}_{\theta_o}$     The transition network parameterised by $\theta_s$ or $\theta_o$, respectively.

$\theta$, $\phi$     All the parameters of the generative model, and the variational distribution, respectively.

$a$, $b$, $c$, $d$     Four hyperparameters involved in the computation of $\omega_t$.

$N(s_\tau, a_\tau)$     The number of times action $a_\tau$ was explored in state $s_\tau$.

$t$, $\gamma$     The present time step, and the discount factor, respectively.

$\zeta$, $\psi$     The precision of the prior over actions, and the precision of the prior preferences.

$\epsilon$, $\hat{\epsilon}$     The random variable used in the re-parameterisation trick, and a sample of epsilon.

$\mathring{\epsilon}$     The probability of selecting a random actions when using the $\mathring{\epsilon}$-greedy algorithm.

$T_{dec}$     A hyperparameter defining the threshold value corresponding to a clear winner during MCTS.

2*$G(\pi)$, $G_\tau(\pi)$     The expected free energy (EFE) of policy $\pi$, and the EFE received at time

step $\tau$ when following policy $\pi$, respectively.

2*$\bar{\bm{G}}_s$, $\bm{G}_s^{\text{aggr}}$, $\bm{G}_s$, $\bm{N}_s$     The average EFE, the aggregated EFE, the EFE, and the number of visits

of a node $s$, respectively.

$\mu_o, \sigma_o, \mu_a, \sigma_a$     The mean and variance vectors predicted by the encoder and policy networks of the $DAI_{FA}$.

$\mu$, $\sigma$     The mean and variance of the Gaussian distribution over $s_t$ predicted by the encoder.

$\mathring{\mu}$, $\mathring{\sigma}$     The mean and variance of the Gaussian distribution over $s_{t+1}$ predicted by the transition.

$\hat{\mu}$, $\hat{\sigma}$     The mean and variance of the Gaussian distribution over $s_{t+1}$ predicted by the encoder.

$\hat{\pi}$     The parameters of the categorical distribution over $a_t$ predicted by the policy network.

$\omega_t$     The top-down attention parameter modulating the precision of the transition mapping.

$[\,\text{condition}\,]$     An indicator function that equals one if the condition is satisfied and zero otherwise.

$\sigma[\,\bigcdot\,]$     The softmax function.

$\text{Cat}(x; \phi_x)$     A categorical distribution over $x$ parameterised by $\phi_x$.

$\text{Bernoulli}(x; \phi_x)$     A Bernoulli distribution over $x$ parameterised by $\phi_x$.

$\MultiBernoulli(x; \phi_x)$     A product of Bernoulli distributions over $x$ parameterised by $\phi_x$.

2*$\mathcal{N}(x; \mu_x, \sigma_x)$     A multivariate Gaussian over $x$ parameterised by a mean vector $\mu_x$, and

a diagonal covariance matrix whose diagonal elements are $\sigma_x$.

$X \overset{i}{\rightarrow} Y$, $X \overset{m}{\rightarrow} Y$     $X$ is fed as input to $Y$, and the mean of the distribution predicted by $X$ is fed as input to $Y$.

$X \overset{s}{\rightarrow} Y$, $X \rightarrow Y$     a sample from the distribution predicted by $X$ is fed as input to $Y$, and $X$ outputs $Y$.

Notation of Sections <ref> and <ref>.

§ APPENDIX B: $DAI_{MC}$ DISCREPANCIES BETWEEN THE PAPER AND THE CODE

In this section, we focus on the authors' implementation of $DAI_{MC}$ available on GitHub: <https://github.com/zfountas/deep-active-inference-mc/>. First, according to a personal communication with one of the authors, the code available on GitHub (on the 6th of June 2022) is not the same as the one used to run the experiments of the paper. Below, we describe the discrepancies between the paper and the code. For example, the computation of $\omega_t$ in the paper is as follows:

\begin{align*}
\omega_t = \frac{a}{1 + \exp(-\frac{b - D_t}{c})} + d,
\end{align*}

while the code uses the following formula:

\begin{align*}
\omega_t = a\times \Bigg(1 - \frac{1}{1 + \exp\big(-\frac{D_t - b}{c}\big)}\Bigg) + d.
\end{align*}

Also, the paper states that MCTS is perfomed to compute the prior over policies during training. However, in the code, MCTS is only used when testing the model, i.e., no MCTS when training the agent. Additionally, the paper states that actions are selected by sampling from:

\begin{align*}
\tilde{P}(a_t) = \frac{N(\hat{s}_t, a_t)}{\sum_{\hat{a}_t}N(\hat{s}_t, \hat{a}_t)}.
\end{align*}

However, the code selects an entire sequence of actions $\pi = (\mathring{a}_t, \mathring{a}_{t+1}, ..., \mathring{a}_{t+n})$ recursively from the root node in the tree. At each step in the recursion, the node with the highest number of visits $\mathring{a}_\tau$ is selected. Then, actions cancelling each other are removed from the sequence, e.g., if $a_\tau = LEFT$ and $a_{\tau+1} = RIGHT$ then both actions are removed from the sequence. This procedure generates a new sequence of actions $\pi'$ of equal or smaller length. Finally, the entire sequence of actions $\pi'$  is performed in the environment. This avoids the repetition of the planning process for each action-perception cycle (saving computational time), however, this also requires domain knowledge (to remove actions that cancel each other out).

Additionally, in the paper, experiments are run on both the dSprites environment and the animal AI environment. However, the code does not allow the replication of the results on the animal AI environment, i.e., the code handling the animal AI environment has been removed. In addition, the evaluation of the expected free energy is non trivial (see below) and the details are not discussed in the paper. Before explaining how the terms of the EFE are computed, we introduce notation that allows us to express those computational steps concisely. For example, we note:

$$o^r_t \overset{i}{\rightarrow} \text{Encoder} \overset{s}{\rightarrow} \text{Transition} \overset{m}{\rightarrow} \mathring{s}^r_{t+1},$$

meaning that $o^r_t$ is used as input ($\overset{i}{\rightarrow}$) for the encoder, then a state is sampled ($\overset{s}{\rightarrow}$) from the distribution predicted by the encoder and used as input for the transition network, finally, the mean ($\overset{m}{\rightarrow}$) of the distribution predicted by the transition network is used as a maximum aposteriori estimate of $\mathring{s}^r_{t+1}$. Note, the transition network takes two inputs (i.e., a state and an action), when using our concise notation we implicitly assume that the actions prescribed by the policy[$\pi$ is the policy for which the expected free energy is being computed.] $\pi$ are provided as input to the transition network. Also, for each time step $\tau$, the reward $r_\tau$ collected by the agent is encoded in the pixels of the image $o_\tau$ as explained in Figure <ref>, leading to a new image $o_\tau^r$. As illustrated on the right of Figure <ref>, the encoder/decoder networks are trained to predict the resulting images $o_\tau^r$. The computation of the first term in equation (<ref>) is illustrated on the left of Figure <ref>. Concisely, we have:

$$o_t^r \overset{i}{\rightarrow} \text{Encoder} \overset{s}{\rightarrow} \text{Transition} \overset{s}{\rightarrow} \text{Decoder} \overset{m}{\rightarrow} \mathring{o}_{t+1}^r.$$

Next, a matrix ($\mathring{r}_{t+1}$) encoding the maximum reward that the agent can gather is used as parameter of Bernoulli distributions to compute the logarithm of the probability (i.e., $\bm{L}$) of the three first rows of the reconstructed image $\mathring{o}_{t+1}^r$. Note, as explained in Figure <ref>, the first three rows contain the predicted reward obtained at time $t+1$. Finally, the mean of $\bm{L}$ is then computed and is multiplied by ten to get $\mathbb{E}_{\tilde{Q}}[\ln \tilde{P}(o_\tau|\pi)]$. Similarly, the computation of $H[Q(s_\tau|\pi)]$ proceeds as follows:

$$o_t^r \overset{i}{\rightarrow} \text{Encoder} \overset{s}{\rightarrow} \text{Transition} \rightarrow \mathring{\mu}, \ln \mathring{\sigma},$$

where $Q(s_\tau|\pi)$ is equated with $\mathcal{N}(s_\tau;\mathring{\mu}, \mathring{\sigma})$, and an analytical solution is used to compute the entropy of $Q(s_\tau|\pi)$. Next, the computation of $H[Q(o_\tau|s_\tau, \pi)]$ proceeds  as follows:

$$o_t^r \overset{i}{\rightarrow} \text{Encoder} \overset{s}{\rightarrow} \text{Transition} \overset{s}{\rightarrow} \text{Decoder} \overset{m}{\rightarrow} \mathring{o}^r_{t+1},$$

where observation $\mathring{o}^r_{t+1}$ is equated to the parameters of the Bernoulli distribution $Q(o_\tau|s_\tau, \pi)$, and an analytical solution is used to compute $H[Q(o_\tau|s_\tau, \pi)]$. Surprisingly, another observation $\mathring{o}^r_{t+1}$ sampled exactly as before is equated to the parameters of the Bernoulli distribution $Q(o_\tau|s_\tau,\theta,\pi)$, and the same analytical solution is used to compute $H[Q(o_\tau|s_\tau,\theta,\pi)]$. Finally, $H[Q(s_\tau|o_\tau, \pi)]$ is computed by feeding $\mathring{o}^r_{t+1}$ back into the encoder to obtain the mean and log-variance of the Gaussian distribution $Q(s_\tau|o_\tau, \pi)$, and the analytical solution for the entropy of a Gaussian is used to compute $H[Q(s_\tau|o_\tau, \pi)]$.

In summary, two samples of $\mathring{o}^r_{t+1}$ (sampled as described in Figure <ref>) have been equated to the parameters of two diferent distributions, i.e., $Q(o_\tau|s_\tau,\theta,\pi)$, and $Q(o_\tau|s_\tau, \pi)$. Additionally, a third sample of $\mathring{o}^r_{t+1}$ (sampled in the same way) has also been used as input to the distribution $\tilde{P}(o_\tau|\pi)$. Lastly, while Fountas et al., 2020 defines the EFE as in (<ref>), the code turns a plus into a minus, leading to the following definition of the EFE:

\begin{align*}
G_{\tau}(\pi) = &- \mathbb{E}_{\tilde{Q}}\Big[ \ln \tilde{P}(o_\tau|\pi)\Big]\nonumber\\
&\,{\color{red}-} \,\, \mathbb{E}_{Q(\theta|\pi)}\Big[ \mathbb{E}_{Q(o_\tau|\theta,\pi)}\big[H[Q(s_\tau|o_\tau, \pi)]\big] - H[Q(s_\tau|\pi)] \Big]\nonumber\\
&+ \mathbb{E}_{Q(\theta|\pi)Q(s_\tau|\theta,\pi)}\Big[H[Q(o_\tau|s_\tau,\theta,\pi)] \Big] - \mathbb{E}_{Q(s_\tau|\pi)}\Big[ H\big[ Q(o_\tau|s_\tau, \pi) \big] \Big],
\end{align*}

where the red minus was a plus.

[square/.style=regular polygon,regular polygon sides=4]
at (0,0) ;

This figure illustrates how the reward $r_\tau \in [-1, 1]$ is encoded in image $o_\tau$. On the left, the plus and minus signs shows where the reward will be encoded in the image if the reward is positive or negative, respectively. In the middle, a positive reward is being encoded on the left side of the image. On the right, a negative reward is being encoded on the right of the image.

[square/.style=regular polygon,regular polygon sides=4]
[fill=gray!20!white] (-1,-0.5) rectangle (1,0.5);
at (0, 0) $\text{env}_t$;
[fill=gray!20!white] (9.5,-0.5) rectangle (11.5,0.5);
at (10.5, 0) $\text{env}_{t+1}$;
at (5.25, 0.5) $\text{perform action } \hat{a}_t$;
[-latex] (1,0) – (9.5,0);
[-latex] (-0.5,-0.5) – (-0.5,-1) – (0.1,-1);
[-latex] (-0.5,-0.5) – (-0.5,-2) – (0.1,-2);

[-latex] (10,-0.5) – (10,-1) – (10.6,-1);
[-latex] (10,-0.5) – (10,-2) – (10.6,-2);

[anchor=west] at (0, -1) $r_t \,\,\, \sim U([-1, 1])$;
[anchor=west] at (0, -2) $o_t$;
[very thick, dashed] (-0.2,-0.8) rectangle (0.6,-2.2);
[very thick, dashed, -latex] (0.2,-2.2) – (0.2,-3.1);

[anchor=west] at (10.5, -1) $r_{t+1} \,\, = \text{env}_{t+1}.\text{compute\_reward}()$;
[anchor=west] at (10.5, -2) $o_{t+1}$;
[very thick, dashed] (10.3,-0.8) rectangle (11.4,-2.2);
[very thick, dashed, -latex] (10.85,-2.2) – (10.85,-3.1);

[very thick, dashed] (10.85,-2.5) – (7.5,-2.5);
[very thick, dashed] (0.2,-2.5) – (2.9,-2.5);
at (5.25,-2.5) encode reward in image;

[rotate=-90,xshift=4cm,yshift=19.85cm, scale=0.9]vae=$o_{t+1}^r$/$\hat{\mu}$/$\ln \hat{\sigma}$/$\hat{s}_{t+1}^r$/$\hat{o}_{t+1}^r$;

[very thick] (12.5,-2.7) – (13,-2.7) – (13,-15.65) – (12.5,-15.65);
[very thick] (13,-9.175) – (13.5,-9.175);
[anchor=west,text width=3cm,align=center] at (13.3,-9.175) The variational auto-encoder at time $t+1$ is used to train the encoder and decoder networks.;

[rotate=-90, xshift=4cm,yshift=9.2cm, scale=0.9]fountas=$o_t^r$/$\mu$/$\ln \sigma$/$\hat{s}_t^r$/$\mathring{o}_{t+1}^r$;

(1.5,-13.9) rectangle (3.5,-14.9);
(-1,-13.9) rectangle (1,-14.9);
(0,-13.9) – (0,-14.9);
[-latex] (4.7,-14.4) – (3.5,-14.4);
at (-1.7,-14.4) $\mathring{r}_{t+1} =$;
at (0.5,-14.4) $\bm{0}$;
at (-0.5,-14.4) $\bm{1}$;
at (2.5,-14.4) $R(\mathring{o}_{t+1}^r)$;

at (1.25,-13.3) $\bm{L} = \ln \mathcal{B}\text{ernoulli}(R(\mathring{o}_{t+1}^r); \mathring{r}_{t+1})$;
at (1.25,-12.3) $\mathbb{E}_{\tilde{Q}}[\ln \tilde{P}(o_\tau|\pi)] = \text{mean}(\bm{L}) \times 10$;

(2.3,-3.5) rectangle (8.75,-7.5);
[anchor=west] at (2.5,-4) def env$_{t+1}$.compute_reward():;
[anchor=west] at (3,-4.5) $r_{t+1} = r_t \times 0.95;$;
[anchor=west] at (3,-5) if agent.y $\geq$ 32:;
[anchor=west] at (3.5,-5.5) if shape == “square":;
[anchor=west] at (4,-6) $r_{t+1} = \frac{15.5 - \text{agent.x}}{16}$;
[anchor=west] at (3.5,-6.5) else:;
[anchor=west] at (4,-7) $r_{t+1} = \frac{\text{agent.x} - 15.5}{16}$;

This figure illustrates the computation of $\mathbb{E}_{\tilde{Q}}[\ln \tilde{P}(o_\tau|\pi)]$ in (the code of) Fountas et al., 2020. The environment at time $t$ provides the agent with an image $o_t$ and a reward $r_t$ randomly sampled from the interval $[-1;1]$. Then, action $\hat{a}_t$ is performed in the environment and the agent observes an image $o_{t+1}$ and a reward $r_{t+1}$, where $r_{t+1}$ is computed according to the function presented in the center of the image. Next, the reward at time $t$ and $t+1$ are encoded in the images received at time $t$ and $t+1$, respectively, c.f., Figure <ref> for details about the encoding. The encoded image at time $t+1$ (i.e., $o^r_{t+1}$) is then fed into the encoder, the re-parameterisation trick is then used to sample a state from the variational posterior. This state is fed into the decoder which tries to reconstruct the image inputed into the encoder. Once $\hat{o}^r_{t+1}$ has been computed, the weights of the encoder and decoder are learned using back-propagation. On the other hand, the encoded image at time $t$ (i.e., $o^r_t$) is used to compute $\mathbb{E}_{\tilde{Q}}[\ln \tilde{P}(o_\tau|\pi)]$. More precisely, the $o^r_t$ is fed into the encoder, and a state is sampled from the variational posterior $Q_{\phi_s}(s_t)$. This state is then fed as input into the transition network along with the action prescribed by $\pi$ at time $t$, i.e., $a_t$. A state at time $t+1$ can then be sampled from the distribution predicted by the transition network. This state is then inputed into the decoder, which outputs $\mathring{o}^r_{t+1}$. Next, a matrix (i.e., $\mathring{r}_{t+1}$) encoding the maximum reward that the agent can gather is used as a parameter of a Bernoulli distribution to compute the logarithm of the probability (i.e., $\bm{L}$) of the first three rows of $\hat{o}^r_{t+1}$, i.e., $R(\mathring{o}^r_{t+1})$. The mean of $\bm{L}$ is then computed and is multiplied by ten to obtain $\mathbb{E}_{\tilde{Q}}[\ln \tilde{P}(o_\tau|\pi)]$.

§ APPENDIX C: ANALYSIS OF THE EPISTEMIC VALUE

In this appendix, we study the expected free energy decomposition into epistemic and extrinsic value as given in Equation (10) of [Parr and Friston, 2019]. A particular reason for being interrested in this formulation is that it is the version of the expected free energy that has the most influenced the deep active inference approaches, such as Fountas et al., 2020. Starting with the definition of the expected free energy, see Equation (<ref>) in the main-body:

\begin{align*}
G_\tau(\pi) &= \mathbb{E}_{P(o_\tau|s_\tau)Q(s_\tau | \pi)}\big[\ln Q(s_\tau | \pi) - \ln P(o_\tau, s_\tau)\big]\nonumber\\
&= \mathbb{E}_{P(o_\tau|s_\tau)Q(s_\tau | \pi)}\big[\ln Q(s_\tau | \pi) - \ln P(s_\tau|o_\tau)- \ln P(o_\tau)\big]\nonumber\\
&= - \underbrace{\mathbb{E}_{P(o_\tau|s_\tau)Q(s_\tau | \pi)}\big[\ln P(s_\tau|o_\tau) - \ln Q(s_\tau | \pi) \big]}_{\text{Epistemic value}} - \underbrace{\mathbb{E}_{P(o_\tau|s_\tau)Q(s_\tau | \pi)}\big[\ln P(o_\tau)\big]}_{\text{Extrinsic value}}\nonumber
\end{align*}

In the rest of this appendix, we will focus on the epistemic value and report two experiments. In the first experiment, the prior over states (c.f., left-most graph in Figure <ref>) is equal to the approximate posterior (c.f., Figure <ref>), and the likelihood (c.f., Figure <ref>) is becoming more and more uniform (c.f., Figure <ref>). Note, if the likelihood becomes uniform, then the true posterior $P(s_\tau | o_\tau)$ becomes equal to the prior over states $P(s_\tau)$, i.e.,

$$P(s_\tau | o_\tau) \propto P(o_\tau | s_\tau ) P(s_\tau) =  \frac{1}{|o_\tau|} P(s_\tau),$$

where $|o_\tau|$ is the number of observations at time step $\tau$, and after renormalisation $P(s_\tau | o_\tau) = P(s_\tau)$. The left-most graph of Figure <ref> shows that the epistemic value decreases as the likelihood becomes more uniform, i.e., the epistemic value decreases as the true posterior $P(s_\tau | o_\tau)$ becomes more similar to the approximate posterior $Q(s_\tau | \pi)$. This behaviour is to be expected, as the epistemic value is bigger when the true posterior $P(s_\tau | o_\tau)$ and the approximate posterior $Q(s_\tau | \pi)$ are more different. Thus, maximising epistemic value will promote exploration and information gain.

In the second experiment, the likelihood has high entropy (c.f., right-most graph in Figure <ref>), and the prior over states is shifting from left to right (c.f., Figure <ref>). When the prior over states $P(s_\tau)$ and the approximate posterior $Q(s_\tau | \pi)$ are different, the joint distribution $P(o_\tau | s_\tau)Q(s_\tau | \pi)$ will be more similar to the approximate posterior $Q(s_\tau | \pi)$ than it is to the true posterior $P(s_\tau | o_\tau)$. Indeed, as the likelihood is almost uniform, the true posterior $P(s_\tau | o_\tau)$ will almost be equal to the prior over states $P(s_\tau)$. At the same time, as the likelihood is almost uniform, it will not have much impact on the joint distribution $P(o_\tau | s_\tau)Q(s_\tau | \pi)$ and the approximate posterior $Q(s_\tau | \pi)$ will dominate. To sum up, when the prior over states $P(s_\tau)$ and the approximate posterior $Q(s_\tau | \pi)$ are different:

* the true posterior will almost be equal to the prior over states $P(s_\tau | o_\tau) \approx P(s_\tau)$

* the joint distribution $P(o_\tau | s_\tau)Q(s_\tau | \pi)$ will be more similar to $Q(s_\tau | \pi)$ than it is to $P(s_\tau | o_\tau)$

Therefore, the joint distribution $P(o_\tau | s_\tau)Q(s_\tau | \pi)$ will tend to be large when the difference within the expectation is negative (c.f., Figure <ref>). Indeed, as the joint is similar to the approximate posterior, it means that the joint distribution is large when the approximate posterior is large and the true posterior is smaller. This implies that the logarithm of true posterior will be very negative, while the logarithm of the approximate posterior will be less negative, i.e., closer to zero. Then, the logarithm of the approximate posterior will be substrated from the logarithm of the true posterior, i.e., a small positive number will be added to a very negative number. Therefore, the result will be negative.

The right-most graph of Figure <ref> shows that the epistemic value increases as the prior $P(s_\tau)$ and therefore the true posterior $P(o_\tau | s_\tau)$ becomes more similar to the approximate posterior $Q(s_\tau | \pi)$. This behaviour should not be observed, as the epistemic value is bigger when the true posterior $P(s_\tau | o_\tau)$ and the approximate posterior $Q(s_\tau | \pi)$ are more similar. Thus, maximising epistemic value will not promote exploration. In fact, it will recommend a strong focus on a single action as was observed in Figure <ref>.

This figure illustrates the approximate posterior $Q(s_\tau | \pi)$, which is distributed according to a binomial distribution corresponding to 6 trials with a probability of success of 0.5.

This figure provides two views of the same likelihood mapping $P(o_\tau | s_\tau)$, i.e., one from the front and one from behind. The likelihood was created by sliding a binomial distribution (corresponding to 6 trials with a probability of success of 0.5) across the observation axis, each time the state increases by one. Finally, when the binomial reached its most extreme position, it stays the same for the remaining states.

During the first experiement, we changed the likelihood mapping $P(o_\tau | s_\tau)$ by making it more and more uniform. This change is shown in the figure from left to right.

During the second experiement, we changed the prior over states $P(s_\tau)$ by making it more and more different to the approximate posterior $Q(s_\tau | \pi)$. This change is shown in the figure from left to right.

To compute the epistemic value, we first computed the difference between the logarithm of the true posterior $\ln P(s_\tau | o_\tau)$ and the logarithm of the approximate posterior $\ln Q(s_\tau | \pi)$ for all the values taken by the observation $o_\tau$ (c.f., left-most figure). Then, we computed the joint distribution $P(o_\tau | s_\tau)Q(s_\tau | \pi)$ used in the expectation (c.f., middle figure). Next, we compute the element-wise product between the matrices illustrated in the left-most and middle figures (c.f., right-most figure). Finally, the epistemic value is obtained by summing up all the elememts of the element-wise product. In this instance, the epistemic value will be strongly negative.