# Dual-path Mamba: Short and Long-term Bidirectional Selective Structured
State Space Models for Speech Separation

###### Abstract

Transformers have been the most successful architecture for various speech
modeling tasks, including speech separation. However, the self-attention
mechanism in transformers with quadratic complexity is inefficient in
computation and memory. Recent models incorporate new layers and modules along
with transformers for better performance but also introduce extra model
complexity. In this work, we replace transformers with Mamba, a selective
state space model, for speech separation. We propose dual-path Mamba, which
models short-term and long-term forward and backward dependency of speech
signals using selective state spaces. Our experimental results on the
WSJ0-2mix data show that our dual-path Mamba models match or outperform dual-
path transformer models Sepformer with only 60% of its parameters, and the
QDPN with only 30% of its parameters. Our large model also reaches a new
state-of-the-art SI-SNRi of 24.4 dB.

Index Terms—  Speech separation, source separation, speech sequence modeling,
state space model, deep learning

##  1 Introduction

Speech is an inherently long signal of thousands of samples. Speech separation
models separate multiple individual speech, each of thousands of samples, from
a single overlapped speech mixture of the same length. Therefore, an effective
mechanism for modeling long speech sequences is crucial to achieving high
performance in speech separation. With the advancement of deep learning, a
variety of neural sequence modeling architectures, including convolutional
neural networks (CNN) [1], recurrent neural networks (RNN) [2], and
transformers [3], have been adopted into speech separation models and achieved
a state-of-the-art (SOTA) performance at the time they were adopted [4, 5, 6].
However, each architecture faces its own challenges: CNNs are restricted by a
finite receptive field; RNNs are hard to parallelize and struggle with
vanishing or exploding gradients [7]; Transformers often outperform CNNs and
RNNs but require quadratic computation in computing self-attention [3].
Although alternative self-attention mechanisms with sub-quadratic complexity
exist, they fail to match the performance of the original attention mechanism
in speech separation [8]. Meanwhile, the attempt to scale up transformer
parameters, as seen from Sepformer [6] to almost eight times larger QDPN [9],
yields a modest performance improvement, suggesting diminishing returns with
increasing model size. Recently, MossFormer2, which integrates transformers
and RNN-free recurrent networks, has achieved better performance but is much
smaller than QDPN, hinting that speech separation might still benefit from
exploring beyond transformer architectures.

State space model (SSM) [10, 11, 12] is another unique class of sequence
modeling architecture. Recently, a novel selective SSM named Mamba is proposed
[13]. Distinct from earlier SSMs, Mamba incorporates an input-dependent
selection mechanism that improves sequence modeling performance but still
enjoys linear complexity with respect to the sequence length. Mamba models
have matched and even surpassed transformers of comparable sizes in sequence
modeling tasks of text, audio, image, and genomics. [13, 14].

Inspired by Mamba’s effectiveness and efficiency in sequence modeling, we
introduce it into speech separation and propose a new model named dual-path
Mamba (DPMamba). As the name suggests, we follow the long sequence modeling
method in dual-path RNN [5] to split a long speech into multiple short chunks
and apply Mamba models within each chunk, across all chunks, in the original
direction, and in the reversed direction of time. Our experiments of DPMamba
of four different sizes on the WSJ0-2mix dataset [15] demonstrate on-par or
superior performance over SOTA models of similar or larger sizes, including
the CNN-based Wavesplit [16], the RNN-based DPRNN [5] and VSUNOS [17], and the
transformer-based SepFormer [6] and QDPN [9]. The largest model outperforms
Mossformer2 [18] and reaches a new SI-SNRi record of 24.4 dB.

Fig. 1: A top-down view of DPMamba from I to IV.

##  2 Related Works

###  2.1 Speech Separation

Following the shift to adaptive front-ends over the short-time Fourier
transform for waveform encoding [19, 20], time-domain speech separation models
have become the mainstream of research. The core of these models is a mask
estimation network which consists of multiple separation blocks of the same
architecture. The architecture of these blocks plays a crucial role in
separation performance and has thus become a primary focus of research
efforts. CNN and RNN models were first introduced, succeeded by better
transformer and hybrid models. We spotlight models that have set benchmarks in
speech separation across different architecture families and will compare them
with our mamba separation models in Section 4.

CNN: Conv-TasNet [4] first outperforms ideal time-frequency masking in speech
separation. It uses temporal convolutional networks (TCN) with progressively
increasing dilation factors to expand its receptive field. Sudo rm -rf [21,
22] employs a more efficient CNN architecture with down-sampling layers to
expand the receptive field. Wavesplit [16] builds on TCNs but introduces
additional speaker embeddings, which leads to a better performance.

RNN: An earliest time-domain model TasNet [20] utilizes long short-term memory
(LSTM) to separate speech. To improve the modeling ability of long sequence,
DPRNN [5] proposes splitting a speech sequence into smaller chunks and
utilizing two LSTMs for intra-chunk and inter-chunk processing. VSUNOS [17]
outperforms DPRNN by substituting LSTMs with gated RNNs.

Transformer: DPTNet [23] and Sepformer [6] adopt the dual-path architecture of
DPRNN and replace LSTMs with transformers. They differ in the stride size and
the use of a RNN or positional encoding to embed the order information of a
speech sequence. QDPN [9] proposes an improved quasi-dual-path architecture
and scales up parameters to 200M. This marks the best performance of the
largest transformer model in WSJ0-2mix.

Hybrid: More recent separation models employ a hybrid architecture which often
augments a vanilla transformer using another model, like convolution-augmented
transformers in TD-Conformer [24], gated transformer with convolution-
augmented attentions in Mossformer [25], and transformer and RNN-free
recurrent network in Mossformer2 [18]. Separate and Diffuse [26] adds a
generative model to post-refine the speeches separated by Sepformer.

###  2.2 Applications of Mamba

Mamba has been proven a transformer-level performance in multiple modalities
and tasks. Mamba is first applied to text, audio, and genomics modeling tasks
in its own paper [13]. Later applications extend to images [14], biomedical
data [27], graphs [28], motions [29], videos [30], and point clouds [31]. For
speech, [32] adopts Mamba for multi-channel speech enhancement. [33] adopts S4
(the predecessor of Mamba) [12] for single-channel speech separation. To our
knowledge, ours is the first work that adopts Mamba in single-channel speech
separation.

##  3 DPMamba

In this section, we will first review the selective state space model which is
the core component of Mamba. Then, we will explain DPMamba in a top-down way
with Fig. 1. We focus on our adoption of Mamba (purple parts in Fig. 1) into
separation since the rest is commonly shared with other dual-path separation
models.

###  3.1 Selective State Space Model

A state space model (SSM) performs a sequence-to-sequence mapping
$\textbf{x}(t)\in\mathbb{R}\mapsto\textbf{y}(t)\in\mathbb{R}$ with a hidden
state $\textbf{h}\in\mathbb{R}^{H}$, a state transition matrix
$\textbf{A}\in\mathbb{R}^{H\times H}$, an input projection matrix
$\textbf{B}\in\mathbb{R}^{H\times 1}$, and an output projection matrix
$\textbf{C}\in\mathbb{R}^{1\times H}$. $H$ is the dimension of the hidden
state:

$\displaystyle\textbf{h}^{\prime}(t)=\textbf{A}\textbf{h}(t)+\textbf{B}\textbf{x}(t),\hskip
4.0pt\textbf{y}(t)=\textbf{C}\textbf{h}(t)$ (1)

To compute Eq.1 for discrete-time signals in digital devices, we need to
discretize the SSM with discretized matrices $\mkern
1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern 1.5mu$ and $\mkern
1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern 1.5mu$:

$\displaystyle\textbf{h}_{t}=\mkern
1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern
1.5mu\textbf{h}_{t-1}+\mkern
1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern
1.5mu\textbf{x}_{t},\hskip 4.0pt\textbf{y}_{t}=\textbf{C}\textbf{h}_{t}$ (2)

$\mkern 1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern 1.5mu$ and
$\mkern 1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern 1.5mu$ are
approximated by zero-order hold. A learnable parameter $\Delta$ balances how
much to focus or ignore the current state and input:

$\displaystyle\mkern 1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern
1.5mu=\text{exp}(\textbf{$\Delta$}\textbf{A}),\hskip 4.0pt\mkern
1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern
1.5mu=(\textbf{$\Delta$}\textbf{A})^{-1}(\text{exp}(\textbf{$\Delta$}\textbf{A})-I)\cdot\textbf{$\Delta$}\textbf{B}$
(3)

By unrolling the computation in Eq.2 along the sequence, it can be observed
that the output sequence $\textbf{y}_{t}$ is the input sequence
$\textbf{x}_{t}$ convolved with a structured kernel $\mkern
1.5mu\overline{\mkern-1.5mu\textbf{K}\mkern-1.5mu}\mkern 1.5mu$ made up by
$\mkern 1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern 1.5mu$,
$\mkern 1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern 1.5mu$, and
C:

$\displaystyle\mkern 1.5mu\overline{\mkern-1.5mu\textbf{K}\mkern-1.5mu}\mkern
1.5mu=(C\mkern 1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern
1.5mu,\mkern
1.5mu\overline{\mkern-1.5mu\textbf{A}\textbf{B}\mkern-1.5mu}\mkern
1.5mu,...,C\mkern 1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern
1.5mu^{L-1}\mkern 1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern
1.5mu),\hskip 4.0pt\textbf{y}=\textbf{x}\ast\mkern
1.5mu\overline{\mkern-1.5mu\textbf{K}\mkern-1.5mu}\mkern 1.5mu$ (4)

The kernel $\mkern 1.5mu\overline{\mkern-1.5mu\textbf{K}\mkern-1.5mu}\mkern
1.5mu$ can be pre-computed if $\mkern
1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern 1.5mu$, $\mkern
1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern 1.5mu$, and C are
fixed. However, the SSM in Mamba is input selective to enable input content-
awareness: It updates the parameters $\textbf{$\Delta$}_{t}$ (and therefore
$\mkern 1.5mu\overline{\mkern-1.5mu\textbf{A}\mkern-1.5mu}\mkern 1.5mu_{t}$),
$\mkern 1.5mu\overline{\mkern-1.5mu\textbf{B}\mkern-1.5mu}\mkern 1.5mu_{t}$,
and $\textbf{C}_{t}$ based on the input $\textbf{x}_{t}$ at each timestep $t$.
To accommodate these dynamic updates, [13] proposes an efficient selective
scan algorithm for the calculation of Eq.4.

###  3.2 Time-domain Dual-path Model

DPMamba follows the time-domain dual-path structure of previous SOTA
separation models [5, 6, 25]. A linear encoder encodes a single-channel wave
mixture $x\in{\mathbb{R}^{1\times T}}$ of $T$ samples to a two-dimensional
latent representation $h_{x}\in{\mathbb{R}^{D\times N}}$ of $D$ dimensions and
$N$ frames. The encoder has a kernel size of $16$ and a stride of $8$,
resulting in $N$= $T/8$. The encoder dimension $D$ is a hyperparameter. Next,
for two speakers in the mixture, a Mamba masking network (Mamba MaskNet)
estimates two masks $m_{1},m_{2}\in{\mathbb{R}^{D\times N}}$. These masks are
then applied through element-wise multiplication with $h_{x}$ to isolate
individual speech sources
$h_{\hat{s}_{1}},h_{\hat{s}_{2}}\in{\mathbb{R}^{D\times N}}$. Finally, a
linear decoder decodes $h_{\hat{s}_{1}},h_{\hat{s}_{2}}$ back to waveforms
$\hat{s}_{1},\hat{s}_{2}\in{\mathbb{R}^{1\times T}}$. Fig. 1 I draws the same
workflow.

Mamba MaskNet is a dual-path network that operates on a three-dimensional
tensor $h_{0}\in{\mathbb{R}^{D\times K\times S}}$. $h_{0}$ is obtained by
Chunking $N$ frames of $h_{x}$ into multiple smaller chunks. $K$=250 is the
chunk size, and $S$ is the resulting number of chunks, with an overlap factor
of $50\%$ between adjacent chunks. Mamba MaskNet processes within- and across-
chunk sequences of much smaller lengths $K$ and $S$ than $N$, which are two
substantially easier sequence modeling tasks [5]. The output of Mamba MaskNet,
$h_{R}\in{\mathbb{R}^{D\times K\times S}}$, is converted back to two-
dimensional by OverlapAdd operation. Mask $m_{1}$ and $m_{2}$ are estimated by
a linear layer followed. Additional linear layers are placed before Chunking
and OverlapAdd operation, as drawn in Fig. 1 II.

Note that we use the same chunk size, kernel size, and stride as other dual-
path models like Sepformer [6] and Mossformer [25, 18] for a fair comparison.

###  3.3 Dual-path Bidirectional Mamba

Mamba MaskNet comprises a stack of $R$ dual-path (DP) blocks iteratively
processing the chunked features $h_{0}$, resulting in $h_{1}$ to $h_{R}$ of
the same shape. Within each DP block, four SSMs process the features in four
different ways: intra-chunk forward, intra-chunk backward, inter-chunk
forward, and inter-chunk backward. An intra-chunk SSM locally processes all
$K$ frames within a chunk; An inter-chunk SSM globally processes all $S$
chunks of the entire signal. A forward SSM processes in the original
direction; A backward SSM processes in the opposite direction of the sequence.
As shown in Fig. 1 III and IV, a DP block contains an intra-chunk and an
inter-chunk unit interleaved, and an intra-chunk or inter-chunk unit each
contains a forward and a backward SSM running in parallel.

####  3.3.1 Local and Global

We adopt computation units of the same structure for local and global
processing as shown in Fig. 1 III. Both intra-chunk and inter-chunk units
contain a normalization layer, a bidirectional Mamba (BiMamba), and a skip
connection. By default, we use RMSNorm [34] as the normalization method due to
its computational efficiency. An intermediate feature
$h_{i}\in\mathbb{R}^{D\times K\times S}$ is first processed by the intra-chunk
unit across $K$ frames in every chunk and then processed by the inter-chunk
unit across $S$ chunks in total. The computation in $i$th DP block can be
expressed as follows:

$\displaystyle h_{i}^{\prime}$ $\displaystyle=h_{i}+\overset{\longleftarrow
K\longrightarrow}{\textit{BiMamba}}\bigl{(}\textit{Norm}(h_{i})\bigr{)}$ (5)
$\displaystyle h_{i+1}$ $\displaystyle=h_{i}^{\prime}+\overset{\longleftarrow
S\longrightarrow}{\textit{BiMamba}}\bigl{(}\textit{Norm}(h_{i}^{\prime})\bigr{)}$
(6)

The intra-chunk duration is around 30 ms with our default configuration for 8
kHz speech. In implementation, we permute the last two dimensions $K$ and $S$
of $h_{i}$ ($h_{i}^{\prime}$) to alter between intra-chunk or inter-chunk
processing.

####  3.3.2 Forward and Backward

The original Mamba [13] uses one SSM to process the input sequence in the
forward direction. However, it has been shown that bidirectional models, which
utilize future context, usually outperform unidirectional models in speech
separation [4]. To make Mamba bidirectional, we borrow the BiMamba design from
[14] to run one forward and one backward SSM in parallel, as shown in Fig. 1
IV. For clarity of discussion, we omit the subscripts used in the previous
sections and denote the input and output of BiMamba as $h$ and $h^{+}$,
respectively.

We start with an input sequence $h\in\mathbb{R}^{D\times L}$, where the
dimension $D$ is the same as the encoder dimension and the length
$L\in\\{K,S\\}$ is either the number of frames per chunk or the number of
chunks. A linear layer projects $h$ to
$h^{\scriptscriptstyle\shortrightarrow}\in\mathbb{R}^{E\times L}$, where
$E=2D$ is the dimension expanded by 2. $h^{\scriptscriptstyle\shortleftarrow}$
is $h^{\scriptscriptstyle\shortrightarrow}$ with $L$ samples flipped.
Meanwhile, another linear layer projects $h$ to $z\in\mathbb{R}^{E\times L}$
in order to later gate SSM outputs:

$\displaystyle h^{\scriptscriptstyle\shortrightarrow}$
$\displaystyle=\textit{Linear}_{i}(h),\hskip
3.0pth^{\scriptscriptstyle\shortleftarrow}=\textit{Flip}(h^{\scriptscriptstyle\shortrightarrow})$
(7) $\displaystyle z$ $\displaystyle=\textit{Linear}_{g}(h)$ (8)

The forward and the backward sequence $h^{\scriptscriptstyle\shortrightarrow}$
and $h^{\scriptscriptstyle\shortleftarrow}$ are then processed by their own
convolution followed by activation layers, resulting in
$\textbf{x}^{\scriptscriptstyle\shortrightarrow}$ and
$\textbf{x}^{\scriptscriptstyle\shortleftarrow}$. The convolution has a kernel
size of 4, and the activation $\sigma$ is the Sigmoid Linear Unit (SiLU) [35,
36]:

$\displaystyle\left\\{\begin{aligned} \hskip
5.0pt\textbf{x}^{\scriptscriptstyle\shortrightarrow}&=\sigma\bigl{(}\textit{Conv1d}_{\scriptscriptstyle\shortrightarrow}(h^{\scriptscriptstyle\shortrightarrow})\bigr{)}\\\
\hskip
5.0pt\textbf{x}^{\scriptscriptstyle\shortleftarrow}&=\sigma\bigl{(}\textit{Conv1d}_{\scriptscriptstyle\shortleftarrow}({h}^{\scriptscriptstyle\shortleftarrow})\bigr{)}\end{aligned}\right.$
(9)

$\textbf{x}^{\scriptscriptstyle\shortrightarrow}$ and
$\textbf{x}^{\scriptscriptstyle\shortleftarrow}$ are the inputs to two SSMs
detailed in Section 3.1. The SSM outputs are gated by $\sigma(z)$:

$\displaystyle\left\\{\begin{aligned} \hskip
5.0pt\textbf{y}^{\scriptscriptstyle\shortrightarrow}&=\sigma(z)\otimes\textit{SSM}_{\scriptscriptstyle\shortrightarrow}(\textbf{x}^{\scriptscriptstyle\shortrightarrow})\\\
\hskip
5.0pt\textbf{y}^{\scriptscriptstyle\shortleftarrow}&=\sigma(z)\otimes\textit{SSM}_{\scriptscriptstyle\shortleftarrow}({\textbf{x}}^{\scriptscriptstyle\shortleftarrow})\end{aligned}\right.$
(10)

Finally, the final output $h^{+}$ is obtained by a linear projection of the
average of the forward and the backward processed sequences
$\textbf{y}^{\scriptscriptstyle\shortrightarrow}$ and
$\textbf{y}^{\scriptscriptstyle\shortleftarrow}$, after the latter flipped
back to the original direction:

$\displaystyle
h^{+}=\textit{Linear}_{o}\bigl{(}\frac{\textbf{y}^{\scriptscriptstyle\shortrightarrow}+\textit{Flip}(\textbf{y}^{\scriptscriptstyle\shortleftarrow})}{2}\bigr{)}$
(11)

##  4 Results

###  4.1 Experiments

We implemented DPMamba of four different sizes as documented in Table 1. We
trained them with scale-invariant signal-to-noise ratio (SI-SNR) [37] and
evaluated them with the improvement of SI-SNR (SI-SNRi) and the improvement of
signal-to-distortion ratio (SDRi) on the WSJ0-2mix dataset [15]. We trained
all the models with a batch size of 1, an Adam optimizer [38], a peak learning
rate of $1.5e^{-4}$, a linear learning rate warmup of 20,000 steps (1 epoch),
and a cosine learning rate decay to $0.1$ from the peak. We trained DPMamba
(XS, S, M) for 100 epochs and DPMamba (L) for 150 epochs. We apply speed
perturbation with a random ratio between 95% and 100% and dynamic mixing as
data augmentation. Note that this training setting is similar to the one in
Sepformer [6] but with a different learning rate schedule and fewer epochs.
All the experiments are conducted in a NVIDIA L40 GPU with mixed (bfloat16)
precision.

Table 1: Hyperparameters of DPMamba XS, S, M, and L. $R\times 2$ means $R$ DP
blocks and each block consists of one intra-chunk and one inter-chunk BiMamba
unit.

Model Dimension $D$ #Layers #Params (M) DPMamba (XS) 128 8 $\times$ 2 2.3
DPMamba (S) 256 8 $\times$ 2 8.1 DPMamba (M) 256 16 $\times$ 2 15.9 DPMamba
(L) 512 16 $\times$ 2 59.8

Table 2: A comparison of DPMamba with previous SOTA separation models from
each architecture family on WSJ0-2mix. n.r. stands for not reported.

Model SI-SNRi (dB) SDRi (dB) #Params (M) Stride CNN Conv-TasNet [4] 15.3 15.6
5.1 8 Sudo rm -rf (B=36) [22] 19.5 n.r. 23.2 10 Wavesplit [16] 22.2 22.3 29 1
RNN TasNet [20] 10.8 11.1 n.r. 20 DPRNN [5] 18.8 19.0 2.6 1 VSUNOS [17] 20.1
20.4 7.5 2 Transformer DPTNet [23] 20.2 20.6 2.6 1 Sepformer [6] 22.3 22.4
25.7 8 QDPN [9] 23.6 n.r. 200 8 Hybrid Mossformer (L) [25] 22.8 n.r. 42.1 8
TD-Conformer-XL [24] 21.2 n.r. 102.7 8 Separate And Diffuse [26] 23.9 n.r.
n.r. 8 Mossformer2 (L) [18] 24.1 n.r. 55.7 8 SSM S4M-tiny [33] 19.4 19.7 1.8 8
S4M [33] 20.5 20.7 3.6 8 DPMamba (XS) 18.9 19.1 2.3 8 DPMamba (S) 21.3 21.4
8.1 8 DPMamba (M) 22.4 22.4 15.9 8 DPMamba (L) 24.4 24.6 59.8 8

###  4.2 Performance Comparison

In Table 2, we compare the performance of our models with the performance of
the SOTA models from different architecture families: CNN, RNN, Transformer,
and Hybrid, which we have introduced in Section 2.1. We also include the
performance of S4M [33], which adopts the S4 model for separation and also
belongs to the SSM family. The numbers in the table correspond to their best
models trained with data augmentation, as reported in their papers.

Our four DPMamba models, spanning from the smallest at 2.3 million parameters
to the largest at 59.8 million parameters, achieve SI-SNRi values between 18.9
dB and 24.4 dB and SDRi values from 20.7 dB to 24.6 dB. Notably, the SI-SNRi
of DPMamba (L) exceeds the literature’s highest SI-SNRi of 24.1 dB set by
Mossformer2 (L), establishing a new state-of-the-art on the WSJ0-2mix dataset.
Furthermore, our other sized models also achieve a high performance with a
relatively smaller model size. For instance, DPMamba (XS), despite having
slightly fewer parameters and a stride of 8, outperforms DPRNN, which uses a
stride of 1. A larger stride means a larger down-sampling ratio by the
encoder, leading to faster inference and smaller memory usage. DPMamba (S) is
marginally larger than VSUNOS but achieves an SI-SNRi over 1 dB higher, with
our stride 4 times larger. DPMamba (M) matches Sepformer’s performance and is
only 0.4 dB lower than Mossformer in SI-SNRi, however, it only requires around
60% and 40% of their parameters, respectively. DPMamba (L), being only 30% the
size of QDPN, achieves a SI-SNRi 0.8 dB higher. Finally, our S, M, and L
models outperform the previous SSM model S4M, although the XS model slightly
lags behind the performance of S4M-tiny. The S4M paper [33] does not implement
models of comparable sizes to our S, M, and L, preventing a conclusive
comparison across all model sizes.

Fig. 2: A comparsion of GPU memory usage of DPMamba with Sepformer and DPRNN.

In Fig. 2, we benchmark the GPU memory consumption during a forward pass of
the models on waveforms of different lengths. We compare our models against
Sepformer and DPRNN with default model configurations implemented by
Speechbrain [39]. DPMamba (M) maintains a Sepformer-level performance but
reduces memory usage by 30%. DPMamba (XS) outperforms DPRNN and consumes only
10% of memory in separating 10-second speeches. The small memory usage marks
the potential of our model in mobile devices.

###  4.3 Ablations

We conducted ablation studies of the DPMamba (S) models. These involved
switching from bidirectional to unidirectional models, decreasing or
increasing the dimension $H$ of the hidden state of SSMs, eliminating dynamic
mixing in training, and replacing RMSNorm with LayerNorm. Our numbers in Table
3 reveal that a variation in the hidden state dimension above 8 and the choice
between RMSNorm and LayerNorm have little impact on the separation
performance. Training with dynamic mixing boosts the SI-SNRi and SDRi by 1.3
dB and 1.2 dB, respectively. The most significant performance improvement, a
boost of over 3 dB, comes from adding backward SSMs in the bidirectional
model, compared to only using forward SSMs in the unidirectional model.

Table 3: Ablations on DPMamba (S).

Configuration SI-SNRi (dB) SDRi (dB) #Params (M) Unidirectional 18.1 18.3 7.4
LayerNorm 21.3 21.5 8.1 H=8 21.3 21.5 7.7 H=32 21.3 21.5 8.9 Without DM 20.0
20.2 8.1 Default DPMamba (S) 21.3 21.4 8.1

##  5 Conclusion

In this work, we introduce DPMamba, a new model for speech separation. DPMamba
utilizes a dual-path network to model local and global aspects of speech
sequences and incorporates bidirectional Mamba blocks for processing the
sequences in forward and backward directions. Our models of four different
sizes either meet or surpass the performance of existing CNN, RNN, and
transformer models of similar or large sizes. DPMamba (L) also sets a new
benchmark on the WSJ0-2mix dataset. Moving forward, we will explore two
directions: enhancing the efficiency of the Mamba separation model and
improving the performance by integrating Mamba with other network layers.

##  6 Acknowledgement

This work was funded by the National Institutes of Health (NIH-NIDCD) and a
grant from Marie-Josee and Henry R. Kravis.

## References

  * [1] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel,  “Backpropagation applied to handwritten zip code recognition,”  Neural Computation, vol. 1, no. 4, pp. 541–551, 1989.
  * [2] Long Short-Term Memory,  “Long short-term memory,”  Neural computation, vol. 9, no. 8, pp. 1735–1780, 2010.
  * [3] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin,  “Attention is all you need,”  Advances in neural information processing systems, vol. 30, 2017\.
  * [4] Yi Luo and Nima Mesgarani,  “Conv-tasnet: Surpassing ideal time–frequency magnitude masking for speech separation,”  IEEE/ACM transactions on audio, speech, and language processing, vol. 27, no. 8, pp. 1256–1266, 2019.
  * [5] Yi Luo, Zhuo Chen, and Takuya Yoshioka,  “Dual-path rnn: efficient long sequence modeling for time-domain single-channel speech separation,”  in ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2020, pp. 46–50.
  * [6] Cem Subakan, Mirco Ravanelli, Samuele Cornell, Mirko Bronzi, and Jianyuan Zhong,  “Attention is all you need in speech separation,”  in ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2021, pp. 21–25.
  * [7] Phong Le and Willem Zuidema,  “Quantifying the vanishing gradient and long distance dependency problem in recursive neural networks and recursive LSTMs,”  in Proceedings of the 1st Workshop on Representation Learning for NLP, Phil Blunsom, Kyunghyun Cho, Shay Cohen, Edward Grefenstette, Karl Moritz Hermann, Laura Rimell, Jason Weston, and Scott Wen-tau Yih, Eds., Berlin, Germany, Aug. 2016, pp. 87–93, Association for Computational Linguistics.
  * [8] Cem Subakan, Mirco Ravanelli, Samuele Cornell, François Grondin, and Mirko Bronzi,  “Exploring self-attention mechanisms for speech separation,”  IEEE/ACM Transactions on Audio, Speech, and Language Processing, vol. 31, pp. 2169–2180, 2022.
  * [9] Joel Rixen and Matthias Renz,  “QDPN - Quasi-dual-path Network for single-channel Speech Separation,”  in Proc. Interspeech 2022, 2022, pp. 5353–5357.
  * [10] Rudolph Emil Kalman,  “A new approach to linear filtering and prediction problems,”  1960\.
  * [11] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré,  “Combining recurrent, convolutional, and continuous-time models with linear state space layers,”  Advances in neural information processing systems, vol. 34, pp. 572–585, 2021.
  * [12] Albert Gu, Karan Goel, and Christopher Ré,  “Efficiently modeling long sequences with structured state spaces,”  in The International Conference on Learning Representations (ICLR), 2022.
  * [13] Albert Gu and Tri Dao,  “Mamba: Linear-time sequence modeling with selective state spaces,”  arXiv preprint arXiv:2312.00752, 2023.
  * [14] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang,  “Vision mamba: Efficient visual representation learning with bidirectional state space model,”  arXiv preprint arXiv:2401.09417, 2024.
  * [15] John R. Hershey, Zhuo Chen, Jonathan Le Roux, and Shinji Watanabe,  “Deep clustering: Discriminative embeddings for segmentation and separation,”  in 2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2016, pp. 31–35.
  * [16] Neil Zeghidour and David Grangier,  “Wavesplit: End-to-end speech separation by speaker clustering,”  IEEE/ACM Transactions on Audio, Speech, and Language Processing, vol. 29, pp. 2840–2849, 2021.
  * [17] Eliya Nachmani, Yossi Adi, and Lior Wolf,  “Voice separation with an unknown number of multiple speakers,”  in International Conference on Machine Learning. PMLR, 2020, pp. 7164–7175.
  * [18] Shengkui Zhao, Yukun Ma, Chongjia Ni, Chong Zhang, Hao Wang, Trung Hieu Nguyen, Kun Zhou, Jiaqi Yip, Dianwen Ng, and Bin Ma,  “Mossformer2: Combining transformer and rnn-free recurrent network for enhanced time-domain monaural speech separation,” 2023.
  * [19] Shrikant Venkataramani, Jonah Casebeer, and Paris Smaragdis,  “End-to-end source separation with adaptive front-ends,”  in 2018 52nd asilomar conference on signals, systems, and computers. IEEE, 2018, pp. 684–688.
  * [20] Yi Luo and Nima Mesgarani,  “Tasnet: time-domain audio separation network for real-time, single-channel speech separation,”  in 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2018, pp. 696–700.
  * [21] Efthymios Tzinis, Zhepei Wang, and Paris Smaragdis,  “Sudo rm-rf: Efficient networks for universal audio source separation,”  in 2020 IEEE 30th International Workshop on Machine Learning for Signal Processing (MLSP). IEEE, 2020, pp. 1–6.
  * [22] Efthymios Tzinis, Zhepei Wang, Xilin Jiang, and Paris Smaragdis,  “Compute and memory efficient universal sound source separation,”  Journal of Signal Processing Systems, vol. 94, no. 2, pp. 245–259, 2022.
  * [23] Jingjing Chen, Qirong Mao, and Dong Liu,  “Dual-Path Transformer Network: Direct Context-Aware Modeling for End-to-End Monaural Speech Separation,”  in Proc. Interspeech 2020, 2020, pp. 2642–2646.
  * [24] William Ravenscroft, Stefan Goetze, and Thomas Hain,  “On time domain conformer models for monaural speech separation in noisy reverberant acoustic environments,”  in 2023 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU). IEEE, 2023, pp. 1–7.
  * [25] Shengkui Zhao and Bin Ma,  “Mossformer: Pushing the performance limit of monaural speech separation using gated single-head transformer with convolution-augmented joint self-attentions,”  in ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2023, pp. 1–5.
  * [26] Shahar Lutati, Eliya Nachmani, and Lior Wolf,  “Separate and diffuse: Using a pretrained diffusion model for better source separation,”  in The Twelfth International Conference on Learning Representations, 2024.
  * [27] Jun Ma, Feifei Li, and Bo Wang,  “U-mamba: Enhancing long-range dependency for biomedical image segmentation,”  arXiv preprint arXiv:2401.04722, 2024.
  * [28] Chloe Wang, Oleksii Tsepa, Jun Ma, and Bo Wang,  “Graph-mamba: Towards long-range graph sequence modeling with selective state spaces,”  arXiv preprint arXiv:2402.00789, 2024.
  * [29] Zeyu Zhang, Akide Liu, Ian Reid, Richard Hartley, Bohan Zhuang, and Hao Tang,  “Motion mamba: Efficient and long sequence motion generation with hierarchical and bidirectional selective ssm,”  arXiv preprint arXiv:2403.07487, 2024.
  * [30] Yijun Yang, Zhaohu Xing, and Lei Zhu,  “Vivim: a video vision mamba for medical video object segmentation,”  arXiv preprint arXiv:2401.14168, 2024.
  * [31] Dingkang Liang, Xin Zhou, Xinyu Wang, Xingkui Zhu, Wei Xu, Zhikang Zou, Xiaoqing Ye, and Xiang Bai,  “Pointmamba: A simple state space model for point cloud analysis,”  arXiv preprint arXiv:2402.10739, 2024.
  * [32] Changsheng Quan and Xiaofei Li,  “Multichannel long-term streaming neural speech enhancement for static and moving speakers,”  arXiv preprint arXiv:2403.07675, 2024.
  * [33] Chen Chen, Chao-Han Huck Yang, Kai Li, Yuchen Hu, Pin-Jui Ku, and Eng Siong Chng,  “A neural state-space model approach to efficient speech separation,” 2023.
  * [34] Biao Zhang and Rico Sennrich,  “Root mean square layer normalization,”  Advances in Neural Information Processing Systems, vol. 32, 2019\.
  * [35] Dan Hendrycks and Kevin Gimpel,  “Gaussian error linear units (gelus),”  arXiv preprint arXiv:1606.08415, 2016.
  * [36] Prajit Ramachandran, Barret Zoph, and Quoc V Le,  “Searching for activation functions,”  arXiv preprint arXiv:1710.05941, 2017.
  * [37] Jonathan Le Roux, Scott Wisdom, Hakan Erdogan, and John R Hershey,  “Sdr–half-baked or well done?,”  in ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2019, pp. 626–630.
  * [38] Diederik P. Kingma and Jimmy 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, Yoshua Bengio and Yann LeCun, Eds., 2015.
  * [39] Mirco Ravanelli, Titouan Parcollet, Peter Plantinga, Aku Rouhe, Samuele Cornell, Loren Lugosch, Cem Subakan, Nauman Dawalatabad, Abdelwahab Heba, Jianyuan Zhong, Ju-Chieh Chou, Sung-Lin Yeh, Szu-Wei Fu, Chien-Feng Liao, Elena Rastorgueva, François Grondin, William Aris, Hwidong Na, Yan Gao, Renato De Mori, and Yoshua Bengio,  “SpeechBrain: A general-purpose speech toolkit,” 2021,  arXiv:2106.04624.


# Deep Inductive Logic Programming meets Reinforcement Learning

Andreas Bueff University of EdinburghUniversity of Edinburgh Vaishak Belle
University of Edinburgh

###### Abstract

One approach to explaining the hierarchical levels of understanding within a
machine learning model is the symbolic method of inductive logic programming
(ILP), which is data efficient and capable of learning first-order logic rules
that can entail data behaviour. A differentiable extension to ILP, so-called
differentiable Neural Logic (dNL) networks, are able to learn Boolean
functions as their neural architecture includes symbolic reasoning. We propose
an application of dNL in the field of Relational Reinforcement Learning (RRL)
to address dynamic continuous environments. This represents an extension of
previous work in applying dNL-based ILP in RRL settings, as our proposed model
updates the architecture to enable it to solve problems in continuous RL
environments. The goal of this research is to improve upon current ILP methods
for use in RRL by incorporating non-linear continuous predicates, allowing RRL
agents to reason and make decisions in dynamic and continuous environments.

##  1 Introduction

One approach to explaining the hierarchical levels of understanding within a
machine learning model is the symbolic method of inductive logic programming
(ILP) [16], which is data efficient and capable of learning first-order logic
rules that can entail data behaviour. Recent contributions to the field have
expanded the ILP framework to allow for end-to-end learning, resulting in
hybrid models that can be classified as neuro-symbolic [7, 17, 23, 4]. The
recent developments in neuro-symbolic ILP have expanded the potential
applications of these models to a wider range of learning challenges,
including reinforcement learning [18]. The ILP-based neuro-symbolic model
implemented in this proposal is a differentiable extension to ILP, so-called
differentiable Neural Logic (dNL) networks [17]. dNL networks are able to
learn Boolean functions as their neural architecture includes symbolic
reasoning. The primary neural layers in the dNL model contain weighted neurons
associated with conjunctions as well as weighted activation neurons associated
with disjunctions, providing a means of logical reasoning on the input as well
as a means of optimisation via gradient descent.

Relational RL (RRL) is concerned with learning policies for decision-making
tasks in complex, discrete relational environments [5]. In RRL, the
environment is characterised by a set of entities and relationships between
them, and the agent’s actions can affect the state of these entities and
relationships. Due to its focus on relational problems, RRL has leveraged the
progress made in ILP but is not limited to one method. Recent efforts have
enabled RRL to effectively tackle more complex RL problems by transitioning
from purely symbolic reasoning to a more neuro-symbolic approach that
leverages neural systems. Earlier research in RRL primarily focused on
planning [5], emphasising model-based learning. Later research, however,
focused on model-free methods that combined neuro-symbolic models and RRL [24]
or on the derivation of interpretable FOL policies [11, 18]. Despite these
advancements, there is limited research in applying RRL to more complex
learning challenges, such as continuous state dynamics. This gap in knowledge
motivates the proposed research.

The work by Payani et al. extended RRL to learn FOL policies using a dNL agent
[18]. Taking the concepts of dNL, Payani et al. combined their dNL-ILP
framework with RRL [18]. The author tested on the block world gaming
environment and took advantage of the declarative bias with provided
background knowledge. The authors expanded RRL to handle complex scene
interpretations and used their dNL-ILP differentiable inductive engine to give
RRL an end-to-end learning framework, so-called dNL-RRL, where the authors
focused on policy gradients in order to improve interpretability and expert
constraints to improve the speed of convergence.

The following paper is an extension that incorporates both a continuous and
non-linear interpretation of [17] and the dNL-RRL agent found in [18]
resulting in a dNL-ILP based agent that can both learn in dynamic continuous
environments typically seen in RL control problems and extract FOL rules which
define the agent policy. Various RL algorithms were explored in evaluating our
dNL-ILP agent and ultimately Soft Actor-Critic was found to perform the best
when solving problems on continuous state spaces. As the contribution bridges
the application of dNL-ILP agents in discrete state spaces with that of
continuous, we evaluated the model on a dynamic RL environment where the
optimal policy captured rules such as those seen with classical mechanics in
physics. In our evaluation, the model was applied to two control problems
including the Cart Pole problem and Lunar Lander problem [2]. From our initial
results, we were able to obtain agent policies which incorporate continuous
predicate functions as well as non-linear continuous predicate functions. We
also add that our proposed agent, while primarily focused on solving classical
RL problems, is referred to as an RRL agent due to its use of relational
language in the derived FOL rules and the ability to incorporate background
knowledge through the use of non-linear predicate functions.

##  2 Background

###  2.1 Reinforcement Learning

Reinforcement learning (RL) aims to find an optimal action sequence for an
agent to maximize its expected future reward. This is typically done by
modelling the environment as a Markov Decision Process (MDP), which is defined
as a tuple
$\mathcal{M}=\langle\mathcal{S},\mathcal{A},\mathcal{R},\mathcal{T},\mathcal{E}\rangle$.
$\mathcal{S}$ is the set of states, $\mathcal{A}$ is the set of actions that
can be taken, $\mathcal{R}(s_{t},a_{t},s_{t+1})$ is the reward function that
takes the current state $s_{t}$, current action $a_{t}$, and returns the
reward from transitioning to the state $s_{t+1}$,
$\mathcal{T}(s_{t},a_{t},s_{t+1})$ is the transition probability function
which represents the probability of transitioning to state $s_{t+1}$ from
state $s_{t}$ given action $a_{t}$ was taken, and
$\mathcal{E}\subset\mathcal{S}$ is the set of terminal states [21, 12]. A
discount factor $\gamma\in[0,1]$ determines the importance of receiving a
reward in the current state versus the future state, where
$R_{t}=\sum_{k=0}^{\infty}\gamma^{k}r(a_{t+k},s_{t+k})$ is the total
accumulated return from the time step. The value function
$V^{\pi}(s)=\mathbb{E}_{\pi}[R_{t}]$ of a policy is the measure of the
expected sum of discounted rewards. The goal is to find the optimal policy
$\pi^{*}=\operatorname*{arg\,max}_{\pi}\mathbb{E}[\sum_{t=0}^{\infty}\gamma^{t}r(a_{t},s_{t})|s_{0}=s]$,
which maps a history of observations to the next action[21, 20, 12, 19, 10,
14]. In practice, at each time step, the agent selects an action based on its
current state and the policy, and receives a reward based on the transition to
the next state. In our investigation of RL methods, we found that Soft Actor-
Critic (SAC) is the best approach for our dNL based agent. SAC combines the
strengths of both value-based and policy-based methods by alternating between
updates to the policy and updates to the value function and Q-function. The
entropy regularisation term is added to the policy update step to encourage
exploration, and the algorithm uses a temperature parameter, $\alpha$, to
control the trade-off between exploration and exploitation. As we evaluate our
agent on environments consisting of discrete actions, we use a discrete action
variant of SAC [3]. The main difference is in the policy update step, where
the objective becomes
$\pi^{*}=\operatorname*{arg\,max}_{\pi}\sum_{t=0}^{T}\mathbb{E}_{s_{t}\sim
D}[\mathbb{E}_{a_{t}\sim\pi}[r(s_{t},a_{t})+\alpha\mathbb{H}(\pi(\cdot|s_{t}))]]$.
Where $\mathbb{H}(\pi(a|s))$ is the entropy of the policy, $\gamma$ is the
discount factor, $D$ is the replay buffer, and $\alpha$ is the temperature
parameter. In the discrete setting, $\pi$ maps states to a vector of
probabilities with $|A|$ elements. For the actor cost function, we use the
following equation $J_{\pi}(\phi)=\mathbb{E}_{s_{t}\sim
D}[\mathbb{E}_{a_{t}\sim\pi}[\alpha\log(\pi_{\phi}(a_{t}|s_{t}))-Q_{\theta}(s_{t},a_{t})]]$.
Where $\pi_{\phi}$ is the policy parameterized by $\phi$ and $Q_{\theta}$ is
the Q-function parameterized by $\theta$. Discrete SAC policy maximises the
probability of discrete actions as opposed to the continuous SAC policy which
optimises two parameters of a Gaussian distribution.

###  2.2 Inductive Logic Programming (ILP)

Inductive logic programming (ILP) is a method of symbolic computing which
automatically constructs logic programs given a background knowledge base (KB)
[16]. An ILP problem is represented as a tuple
$(\mathcal{B},\mathcal{P},\mathcal{N})$ of ground atoms, with $\mathcal{B}$
being the background assumptions, $\mathcal{P}$ being a set of positive
instances which help define the target predicate to be learned, and
$\mathcal{N}$ being the set of negative instances of the target predicate. The
goal of ILP is to construct a logic program that explains all provided
positive sets and rejects the negative ones. Given an ILP problem
$(\mathcal{B},\mathcal{P},\mathcal{N})$, the goal is to identify a set of
hypotheses (clauses) $\mathcal{H}$ such that
$\mathcal{B}\wedge\mathcal{H}\models\gamma$ for all $\gamma\in\mathcal{P}$ and
$\mathcal{B}\wedge\mathcal{H}\not\models\gamma$ for all
$\gamma\in\mathcal{N}$, where $\models$ denotes logical entailment. In other
words, the conjunction of the background knowledge and hypothesis should
entail all positive instances and not entail any negative instances.

###  2.3 Differentiable Neural Logic (dNL)

The core component of the dNL network is their use of differentiable neural
logic layers to learn Boolean functions [17]. The dNL architecture uses
membership weights and conjunctive and disjunctive layers to learn a target
predicate or Boolean function. Learning a target predicate $p$ requires the
construction of Boolean function $\mathcal{F}_{p}$ which passes in a Boolean
vector x of size $N$ with elements $x^{(i)}$, into a neural conjunction
function $f_{conj}$ (see equation 1a) which is defined by a conjunction
Boolean function $F_{conj}$ (see equation 1b). A predicate defined by Boolean
function in this matter is extracted by parsing the architecture for
membership weights ($w^{(i)}$) above a given threshold, where membership
weights are converted to Boolean weights via a sigmoid
$m^{(i)}=\sigma(cw^{(i)})$ with constant $c\geq 1$. Membership weights are
paired with continuous lower and upper bound predicate functions (see equation
3) which are eventually interpreted as atoms in the body of the predicate
being learned. These same Boolean predicate functions are used to transform
non-Boolean data into a Boolean format for the logic layers.

$f_{conj}(\textbf{x})=\prod_{i=1}^{N}F_{conj}(x^{(i)},m^{(i)})$ (1a)

$F_{conj}(x^{(i)},m^{(i)})=\overline{\overline{x^{(i)}}m^{(i)}}=1-m^{(i)}(1-x^{(i)})$
(1b)

Following induction by the conjunctive layer, outputs are fed into a neural
disjunction function $f_{disj}$ (see equation 2a) which is constructed from
disjunctive Boolean functions $F_{disj}$ (see equation 2b). The disjunctive
layer provides multiple definitions for our target predicate if necessary.

$f_{disj}(\textbf{x})=1-\prod_{i=1}^{N}(1-F_{disj}(x^{(i)},m^{(i)}))$ (2a)

$F_{disj}(x^{(i)},m^{(i)})=x^{(i)}m^{(i)}$ (2b)

As mentioned, a target predicate function $\mathcal{F}p$ is defined by the
cascading architecture of the dNL network, and within the network bounded
continuous Boolean predicates associated with a continuous feature are used to
define our target predicate. Boolean predicate functions are used to handle
continuous input by partitioning input into a series of lower and upper-bound
predicates, referred to as continuous predicates as they capture the interval
bounds of continuous features. Similarly, Boolean predicates mapped to
discrete features are referred to as discrete predicates. These bounded
continuous predicates return either true or false when a continuous value
meets the condition. A continuous input $x$ is associated with $k$ pairs of
upper and lower boundary predicates, each pair corresponding to a bounded
range $(l_{xi}<x<u_{xi})$, where $i\in{1,2,\cdots,k}$. A Boolean upper
boundary predicate $gt_{x}^{i}(x,l{xi})$, states whether “x is greater than
$l_{xi}$” is true, and a Boolean lower boundary predicate
$lt_{x}^{i}(x,u_{xi})$ states whether “x is less than $u_{xi}$” is true (see
equation 3). The lower and upper boundary values $l_{xi}$ and $u_{xi}$ are
treated as trainable weights which can be optimised during training. To
clarify the need for both lower and upper predicate bounds, the regions are
complete and disjoint, meaning that each input value belongs to exactly one
region. Although most of the upper and lower bounds may coincide, having both
upper and lower predicates allows for greater flexibility in defining the
regions and ensures that no input values fall outside the defined regions. A
diagram illustrating this partitioning of the input domain can be found in
Figure 1.

$\mathcal{F}_{gt^{i}_{x}}=\sigma(c(x-u_{xi})),\quad\mathcal{F}_{lt^{i}_{x}}=\sigma(-c(x-l_{xi}))$
(3)

$x^{*}=0.8$$\mathcal{F}_{lt_{x}^{1}}(x,0.5)$$\mathcal{F}_{lt_{x}^{2}}(x,1.0)$$\mathcal{F}_{gt_{x}^{1}}(x,0.5)$$\mathcal{F}_{gt_{x}^{2}}(x,1.0)$
Figure 1: Given $k=2$ and feature instance $x^{*}$ with value $0.8$, would
result in two distinct activations of bounded Boolean predicates. Stating
$x^{*}$ is less than $1.0$ and $x^{*}$ is greater than $0.5$.

##  3 Methodology

Developing a neuro-symbolic RL agent that incorporates continuous predicate
functions posed several non-trivial challenges. Firstly, the work of Payani et
al. [17, 18], which our model extends, did not explore non-linear predicate
functions, so we had to develop new functions to capture more complex
relationships in the data (see equation 7). Additionally, while Payani et al.
did introduce continuous predicates in a supervised setting, they only tested
dNL-ILP on two datasets (Wine and Sonar datasets [17, 6]), whereas we made the
integration of continuous predicates a focus of our work. 111Moreover, from an
engineering standpoint, the dNL code was originally written in TensorFlow 1,
which had to be converted to TensorFlow 2 for use in the RL environment.
Another challenge was the fact that the original dNL-RRL model explicitly
dealt with discrete predicates, so we had to adapt the model to incorporate
continuous predicates. Finally, while Payani et al. used the REINFORCE RL
algorithm [22], we had to test out various other RL algorithms, including SAC
[9], to find the most effective implementation for our model.222Extensive
coding was required to convert the SAC algorithm, which typically has a
discrete action space, to TensorFlow 2, as all available examples online were
coded in PyTorch.

In the dNL-RL interpretation, the Boolean target predicate function
$\mathcal{F}_{p}$ is referred to as a discrete action predicate function. The
actions of the agent are the target predicates and so must remain discrete.
For example, an agent that can take an action $turnLeft()$, has the
corresponding discrete action predicate function $\mathcal{F}_{turnLeft}$.
Each discrete action predicate function is defined by an input matrix
$\mathbf{I}$, which is composed of continuous lower and upper bound predicate
functions as well as discrete Boolean predicate functions, all with associated
weights. The set of the continuous state features $cnt_{p}$ and the set of the
discrete state features $dsc_{p}$ for the associated RL environment are given
Boolean interpretations after being passed through the input matrix, as seen
in equation 4. The row dimensions of $\mathbf{I}$ are defined by the batch
size $\mathbf{b}$ and the column dimensions are defined by $N_{e}=(2\times
k\times|cnt_{p}|)+(2\times|dsc_{p}|)$ (assuming discrete features are
Boolean). The state feature ranges used to define the Boolean predicate
functions are determined by equal-width binning, and we define the batch size
as $\mathbf{b}$. The $\mathcal{F}_{{e}}$ function takes in a one-hot encoded
Boolean discrete input.

$\mathbf{I}=\begin{bmatrix}\mathcal{F}_{gt^{1}_{e_{1}}},\mathcal{F}_{lt^{1}_{e_{1}}}&\cdots&\mathcal{F}_{gt^{k}_{e_{1}}},\mathcal{F}_{lt^{k}_{e_{1}}}&\cdots&\mathcal{F}_{gt^{k}_{e_{|cnt_{p}|}}},\mathcal{F}_{lt^{k}_{e_{|cnt_{p}|}}}\\\
\vdots&\vdots&\vdots&\vdots&\vdots\end{bmatrix}\bigoplus\begin{bmatrix}\mathcal{F}_{{e_{1}}},\mathcal{F}_{\overline{e_{1}}}&\cdots&\mathcal{F}_{{e_{|dsc_{p}|}}},\mathcal{F}_{\overline{e_{|dsc_{p}|}}}\\\
\vdots&\vdots&\vdots\ \par\end{bmatrix}$ (4)

In defining the discrete action predicate function $\mathcal{F}_{p}$ for an
action $p$, the function takes as input the input matrix $\mathbf{I}$ as well
as the disjunction layer $F_{disj}$ and the conjunction layer $F_{conj}$ where
$N_{e}$ is used to define the number of layers in the conjunction layer and
$N_{p}$ is used to define the number of layers in the disjunction layer, seen
in equation 5. As an agent learns rules for multiple actions, we define the
set of discrete action predicates as the predicate action policy
$\pi_{\mathcal{F}}$ which is the set
$\\{\mathcal{F}_{p^{1}},\mathcal{F}_{p^{2}},\cdots,\mathcal{F}_{p^{n}}\\}$ for
all actions in the action space $p\in\mathcal{A}$ for the RL environment.

To prevent the gradient from becoming excessively small during training,
initial weights $\mathbf{W}$ are initialised using a negative mean random
Gaussian distribution, which ensures that they are close to zero. To keep
membership weights m between 0 and 1, a constant $c\geq 1$ is applied,
followed by a sigmoid function. In equations 6a and 6b, we apply the
conjunction function to the input matrix using the corresponding membership
weights. Note that $\textbf{W}^{conj}$ is a matrix with dimensions $N_{p}$ and
$N_{e}$ while $\textbf{W}^{disj}$ is a vector of size $N_{p}$. The disjunction
function takes the output of the conjunction function as input as we cascade
our neural architecture.

$\qquad\quad\mathcal{F}_{p}|_{\mathbf{I},N_{e},N_{p}}=F_{disj}(N_{p},F_{conj}(N_{e},\mathbf{I}))$
(5)

$\quad{F}_{conj}=\sum_{i=1}^{N_{e}}[1-m^{conj}_{i}(1-\mathbf{I}_{i})]\quad\text{where}\quad\mathbf{m}^{conj}=\sigma(c\mathbf{W}^{conj}),\
\mathbf{W}^{conj}_{N_{p},N_{e}}\sim\mathcal{N}$ (6a)

$\quad{F}_{disj}=1-\sum_{i=1}^{N_{p}}[1-m_{i}^{disj}{F}_{conj}^{i}]\quad\text{where}\quad\mathbf{m}^{disj}=\sigma(c\mathbf{W}^{disj}),\
\mathbf{W}^{disj}_{1,N_{p}}\sim\mathcal{N}$ (6b)

We incorporate non-linear transformations on some of the state features. As in
the case of various control problems in RL, the calculation of state
transitions is based on non-linear transformations from the current state. We
define the following $k$ upper-boundary and lower-boundary functions as non-
linear transformation predicates in equation 7.

$\mathcal{F}_{gt^{i}_{f(x)}}=\sigma(c(f(x)-u_{xi})),\>\mathcal{F}_{lt^{i}_{f(x)}}=\sigma(-c(f(x)-l_{xi}))\>\text{where}\>f(x)\in\\{sqr(x),sin(x),\cdots\\}$
(7)

Our dNL based agents can be differentiated by the type of prior information
used. The baseline dNL agent uses only continuous and discrete Boolean
predicates and is referred to as a dNL RL continuous (dNLRLc) agent. It learns
the continuous state features without using any prior information from a
knowledge base $\mathsf{KB}_{T}$ which contains transformation functions
mapped to specific continuous state features. However, when prior information
such as non-linear equations or transformations between features is included,
the agent becomes a dNL RL non-linear continuous (dNLRLnlc) agent. This
involves adding non-linear continuous Boolean predicates to the input matrix
as new predicate functions in $\pi_{\mathcal{F}}$. In the experiment section,
we specify which non-linear transformations were used for the dNLRLnlc agent
in each RL environment.

In Algorithm 1 we provide a pseudo-code implementation for a dNL-RL agent. The
agent takes in a state $\mathbf{s}_{t}$, at time step $t$ and returns an
action $\mathbf{a}_{t}$ for that time step. The algorithm iterates through
each feature of the state, $s_{t}^{(i)}$, and separates the elements into two
cases: discrete and continuous. If the element is discrete, it adds the
mapping of the discrete predicate $dsc_{p}[i]$ to the state element
$s_{t}^{(i)}$ in a processed state set $\mathbb{S}_{t}$. If the element is
continuous, it checks if a transformation function for the feature exists in
the knowledge base $\mathsf{KB}_{T}[i]$. If it does, it adds the mapping of
the continuous predicate $cnt_{p}[i]$ to the function applied to the state
feature $s_{t}^{(i)}$ in the state set $\mathbb{S}_{t}$. If the knowledge base
does not exist, it adds the mapping of the continuous predicate $cnt_{p}[i]$
to the state feature $s_{t}^{(i)}$ instead. After iterating through all state
features, the policy is returned by applying the dNLRL policy function to the
set $\mathbb{S}_{t}$, and then sampling an action $\textbf{a}_{t}$ from the
predicate action policy.

Algorithm 1 dNL-RL agent

1:Input: $\mathbf{s}_{t}$-state at time step $t$

2:Output: $\mathbf{a}_{t}$-action at time step $t$

3:for $s_{t}^{(i)}\in\mathbf{s}_{t}$ do

4: if $s_{t}^{(i)}$ is discrete then

5: $\mathbb{S}_{t}\cup\\{dsc_{p}[i]\mapsto s_{t}^{(i)}\\}$

6: end if

7: if $s_{t}^{(i)}$ is continuous then

8: if $\mathsf{KB}_{T}[i]$ exists then

9: $\mathbb{S}_{t}\cup\\{cnt_{p}[i]\mapsto\mathsf{KB}_{T}[i](s_{t}^{(i)})\\}$

10: else

11: $\mathbb{S}_{t}\cup\\{cnt_{p}[i]\mapsto s_{t}^{(i)}\\}$

12: end if

13: end if

14:end for

15:$\mathbb{\pi}_{\mathcal{F}}\leftarrow\mathrm{dNLRL}(\mathbb{S}_{t})$

16:$\mathbf{a}_{t}\sim\mathbb{\pi}_{\mathcal{F}}(\mathcal{F}_{a_{t}}|\mathbb{S}_{t})$

To derive the predicate action policy $\mathbb{\pi}_{\mathcal{F}}$, we perform
Boolean reasoning on the set of processed state features with associated
predicates as seen in Algorithm 2. Here we take as input a set of discrete
predicates $dsc_{p}$, a set of continuous predicates $cnt_{p}$, a set of
target action predicates $\mathbb{P}$, and a set of processed state feature
predicates $\mathbb{S}_{t}$. We loop through each target action predicate in
the set $\mathbb{P}$ and for each predicate, it creates two empty sets used
for constructing the input matrix, one for the continuous predicates
$\mathbf{I}_{c}$ and one for the discrete predicates $\mathbf{I}_{d}$. The
dNLRL policy then loops through each continuous predicate in $cnt_{p}$ and for
each predicate, it creates a disjunction of predicates, which are added to the
continuous input matrix set $\mathbf{I}_{c}$. Then it does the same for each
discrete predicate in $dsc_{p}$ but instead, the discrete predicates are added
to the discrete input matrix set $\mathbf{I}_{d}$. We take the union of these
two sets to get the final input matrix $\mathbf{I}$. Finally, it creates an
evaluated target action predicate $\mathcal{X}_{p}$ by reasoning on target
predicate function $\mathcal{F}_{p}$ as defined in equation 5. This target
action predicate $\mathcal{X}_{p}$ is then added to the predicate action
policy $\mathbb{\pi}_{\mathcal{F}}$.

Algorithm 2 Single Step Forward Chain Model/ dNLRL policy

1:Input: $dsc_{p}$-set of discrete predicates, $cnt_{p}$-set of continuous
predicates, $\mathbb{P}$-set of target action predicates, $\mathbb{S}_{t}$-set
of state feature predicates at time step $t$

2:Output: $\mathbb{\pi}_{\mathcal{F}}$-predicate action policy

3:for $p\in\mathbb{P}$ do

4: $\mathbf{I}_{c}$=[]

5: $\mathbf{I}_{d}$=[]

6: for $e\in cnt_{p}$ do

7:
$\mathbf{I}_{c}\cup\bigvee^{k}_{i=1}\Big{(}\mathcal{F}_{gt^{i}_{e}}(\mathbb{S}_{t}[e])\vee\mathcal{F}_{lt^{i}_{e}}(\mathbb{S}_{t}[e])\Big{)}$

8: end for

9: for $e\in dsc_{p}$ do

10:
$\mathbf{I}_{d}\cup\Big{(}\mathcal{F}_{e}(\mathbb{S}_{t}[e])\vee\mathcal{F}_{\overline{e}}(\mathbb{S}_{t}[e])\Big{)}$

11: end for

12: $\mathbf{I}\leftarrow\mathbf{I}_{d}\cup\mathbf{I}_{c}$

13: $\mathcal{X}_{p}\leftarrow\mathcal{F}_{p}|_{\mathbf{I},N_{e},N_{p}}$

14: $\mathbb{\pi}_{\mathcal{F}}\cup(\mathcal{X}_{p})$

15:end for

##  4 Experiments and Results

###  4.1 Environments and Tasks

The following RL environments are continuous control problems developed by
Open AI Gym [2]. The Cart Pole problem is a benchmark for evaluating RL
algorithms [13]. The goal of the problem is to balance a pole in the upright
position by selecting between two discrete actions (move left, move right).
The state space is 4-dimensional, all continuous. The features are the
following $x:$ Cart Position, $x^{\prime}:$ Cart Velocity, $\theta:$ Pole
Angle, and $\theta^{\prime}:$ Pole Angular Velocity. In the Lunar Lander
environment, an agent learns a policy for throttling the side and main engines
to control the descent of a lander such that it lands on a landing pad. The
discrete actions include activating the right thruster engine, activating the
left thruster engine, activating the main engine, and doing nothing. The state
space is 8-dimensional, 6 inputs are continuous and 2 are discrete. The
continuous state inputs include including $x:$ Lander Position on the X-axis,
$y:$ Lander Position on the Y-axis, $v_{x}:$ Horizontal velocity, $v_{y}:$
Vertical velocity, $\theta:$ Angular orientation in space, and $v_{\theta}:$
Angular velocity. The remaining two features are Boolean, Left: which
indicates if the left leg is touching the ground and Right: which indicates if
the right leg is in contact with the ground.

####  4.1.1 Baselines

During our investigation into the development of our continuous logic-based
algorithm, we sought first to see if the dNLRLc agent could in fact learn
interpretable policies on continuous environments. As stated this led to
evaluation of the dNLRLc architectures on various RL algorithms. This includes
the policy gradient algorithm REINFORCE [22], the model-free algorithm Deep
Q-Network (DQN) [15], the on-policy method of Advantage Actor-Critic (A2C)
[14], and the Soft Actor-Critic (SAC) [9]. In the case of SAC, we use a
variant designed to handle discrete actions [3]. In cases where we inject
prior knowledge in the form of non-linear functions added to the
transformation knowledge base, we designate the model as (SAC-NL).

###  4.2 Algorithm Performance Comparison

We present algorithmic comparisons of the various RL algorithms combined with
the dNLRLc agent where evaluation is carried out on the Cart Pole problem. For
the evaluation, only SAC is used to test the dNLRLnlc agent, as it was the
best-performing algorithm with the dNLRLc agent. In the Lunar Lander
environment, we evaluate a dNLNLc agent and dNLRLnlc agent using SAC entirely.
Each algorithmic framework used the same binning scheme for the continuous
features, that being ‘equal-width binning’ where bins for each feature were
predefined for the learning environment and kept constant between the various
RL algorithms.

####  4.2.1 Task: Cart Pole Problem

For the Cart Pole problem, each algorithmic framework used the same binning
scheme for ‘equal-width binning’ with the number of bins set to $[4,4,4,4]$
for each feature respectively. In Figure 2, we can observe the results
comparing the performances of each dNLRLc agent and paired RL architecture in
the Cart Pole Problem domain. As can be observed, the best performing model is
that of the SAC and SAC-NL. Performance of dNL based agents is noticeably poor
in the cases of the REINFORCE algorithm where rewards cap at $22.0$. DQN and
A2C perform better but rewards plateau far below what is to be expected. The
non-linear addition for the dNLRLnlc agent is a simple transformation of the
Pole Angle $\theta$, as the calculation of the next state $s_{t+1}$ within the
code for the Cart Pole problem is dependent on the $sine$ transformation of
$\theta$. For SAC-NL, we apply a $sine$ transformation and so policies can
contain atoms designated $PoleAngleSine$ with corresponding inequalities. The
performance of the dNLRLnlc agents is not significantly better than the dNLRLc
agent when using SAC. Both agent variants are able to achieve the maximum
reward of $300.0$ consistently with periodic fluctuations. In figure 4.2.1, we
can observe the moving standard deviation for the dNLRLc and dNLRLnlc when
using SAC. Although the fluctuations in the standard deviations are not in
sync, the variation between the two agents is not significant enough to
conclude that one agent is more stable.

Figure 2: Episode scores for RL algorithms relying on dNLRLc and dNLRLnlc
agents.

Figure 4.2.1 presents a comparison of the performance of the various RL
algorithms using standard neural networks in solving the Cart Pole problem.
The figure shows that, while SAC is able to solve the problem and stabilise,
the neural network agents learn faster and are more stable than dNLRLc agents.
Both agents using A2C and DQN exhibit inconsistent performance with rewards
fluctuating greatly and both the neural networks and dNLRLc agents perform
poorly when using REINFORCE. These findings suggest that while dNLRLc models
may offer greater interpretability, they may come at the cost of decreased
performance in some RL tasks. Given that, it is also evident that SAC is the
primary contributor to solving the control problem as even standard neural
network based agents failed otherwise. We also note that the black-box neural
network-based SAC agent outperforms the interpretable dNL-RL SAC agent in
terms of speed, achieving an average convergence time of 2168.31 seconds,
compared to the average convergence time of 10807.94 seconds for the dNL-RL
agent, highlighting a trade-off between interpretability and learning
efficiency, which could be addressed in future research by optimising the
algorithm for faster convergence.

figure Episode scores for RL algorithms relying on neural network based
agents.

figureMoving (window size 20) standard deviation for SAC relying on dNLRLc and
dNLRLnlc agents.

The main advantage of using dNL agents is that returned policies are
interpretable. These policies provide predicate logic interpretations for the
individual discrete actions, in the case of the Cart Pole problem, we have
rules for when to move left() and when to move right(). The policy in our case
is defined as predicate action policy, that is individual actions are treated
as target predicates where the clausal body states what inequalities should be
true for a state feature in order for the action to be considered by the
agent. We also clarify that the clausal body is defined by a conjunction of
atoms where the associated membership weight is placed before the atom in
brackets, extracted from the conjunction neuron. If the membership weight is
above $0.95$, i.e. the model is confident it should be in the definition, then
the weight is not listed next to the atom. Similarly, for the disjunction
neurons, the value in brackets before the rule represents the membership
weights.

CartPole | Policy rules for dNLRLc
---|---
| mean reward:
---
$290.3\pm 32.3$
| left()
---
$:-([0.56]CartPos<2.83\wedge[0.64]CartVeloc>-1.19\wedge CartVeloc>0.18$
$\wedge PoleAngle>-0.62\wedge PoleAngleVeloc>-0.24)$
$:-([0.81]CartPos<2.82\wedge PoleAngle>-0.06\wedge PoleAngleVeloc>0.08)$
right()
$:-(CartVeloc>-1.19\wedge PoleAngle<-0.03\wedge
PoleAngleVeloc<0.34\wedge[0.56]PoleAngleVeloc<2.48)$
$:-(CartPos<0.16\wedge[0.52]CartVeloc<2.11\wedge[0.50]PoleAngle>-0.56$
$\wedge PoleAngle<0.11\wedge PoleAngleVeloc<0.25)$
Table 1: The FOL policy rules for the dNLRLc agent, trained using SAC.
Extracted continuous inequality predicates define rules for each action.

In table 1, the predicate action policy for the dNLRLc agent trained using SAC
is provided. In the body of the discrete action predicate left(), it is
observed that two definitions are returned wherein which we see both
associated disjunction neuron membership weights pass the threshold. The three
state features present in the second definition are Cart Position
($CartPos<2.82$), Pole Angle ($PoleAngle>-0.06$), and Pole Angular Velocity
($PoleAngleVeloc>0.08$). Similarly, the body for the discrete action predicate
right() contains two predicate definitions. In both cases, the agent is
confident in either definition as indicated by the absence of a membership
weight for the corresponding disjunction neuron.

CartPole | Policy rules for dNLRLnlc
---|---
| mean reward:
---
$294.7\pm 25.8$
| left()
---
$:-([0.60]CartPos<2.57\wedge PoleAngleSine>0.00\wedge PoleAngleVeloc>-0.38)$
right()
$:-(PoleAngleSine<0.04\wedge PoleAngleVeloc<0.00)$
$:-(CartPos<0.74\wedge[0.55]CartVeloc>-1.64\wedge CartVeloc<-0.11$
$\qquad\wedge PoleAngleSine<0.65\wedge[0.66]PoleAngleVeloc>-2.04\wedge
PoleAngleVeloc<0.28)$
Table 2: The FOL policy rules for the dNLRLnlc agent, trained using SAC.
Extracted continuous inequality predicates define rules for each action with
the inclusion of non-linear continuous inequality predicates.

In table 2, the predicate action policy for the dNLRLnlc agent trained using
SAC is provided. The definitions of the predicate action policy are comparable
to those given by the dNLRLnlc agent. Note that the discrete action predicate
left() is defined by a single predicate rule. For both discrete action
predicates, we see the presence of the non-linear continuous predicate in the
definitions, that is $(PoleAngleSine>0.00)$ of the discrete action predicate
$left()$ and $(PoleAngleSine<0.00)$ for the first definition of $right()$ and
$(PoleAngleSine<0.65\wedge PoleAngleSine<0.28)$ for the second definition. In
figure 1 and figure 2 the mean reward is also provided with the standard
deviation. The mean reward in this case refers to the mean of the last 100
episodes.

####  4.2.2 Task: Lunar Lander

At present, the Lunar Lander stands as a more challenging environment for the
dNL-RL agents. As the previous control problem environment has demonstrated
that SAC is the best performing algorithm, only SAC was tested on the Lunar
Lander environments. In setting up the experiments, it was found that both the
discretisation scheme and the initial high and low values for the state
features had a significant impact. In the context of the Lunar Lander, a
reward of $200$ indicates the lander successfully landed on the platform. The
current binning scheme being deployed is $[3,3,3,3,3,3,-,-]$ for each feature,
where ($-$) indicates no binning performed as it is a discrete/Boolean
feature. In order to evaluate the consistency and performance of the dNLRLc
and dNLRLnlc agents, multiple trials were conducted using different random
seeds. We ran each experiment 5 times and the resulting policies were
analysed. While only one representative policy is presented for each agent,
any variations or discrepancies across the trials are discussed. The non-
linear change was again a $sine$ transformation on the state feature angular
orientation in space $\theta$ associated with the transformation predicate
AngleSine.

In Figure 3, we can observe the results comparing the performance of the
dNLRLc agent where dark purple corresponds to the mean reward. The results
indicate the agent does in fact learn, however, it is challenging for the dNL
agent to successfully land the lander consistently. The episode rewards show
the agent can land the lander but the mean rewards skew below a true success
of $200$. We note the mean reward across all trials fluctuates around rewards
of a hundred with significant deviations from the mean. In some initial
states, the learned policy does result in successful landings, often over 200.
The rules in the predicate action policy for the Lunar Lander environment are
given in table 3. As the Lunar Lander has four actions, we find that all
disjunction layers pass the parameter threshold and only in a few cases do we
have conjunction weights for specific atoms. For example for the rule
$fireLeft$, the first rule contains a weighted atom $[0.65]CoordX>-1.52$. For
the action predicate, fireRight() the agent learned four rules, where only the
second rules contain listed conjunction layer weights for individual atoms
$[0.60]CoordX>-1.53$ and $[0.94]AngularVeloc>-2.58$. We observe as well, that
while the bounded atom $CoordX$ appears in the definitions for each action,
each is associated with a conjunction weight. Indicating the agent was less
confident about the inclusion of the $CoordX$ continuous predicate function.

Figure 3: Episode scores for SAC algorithms relying on dNLRLc. Dark purple corresponds to the mean reward. Light purple corresponds to the standard deviation and the red line is the moving average for 50 episodes. LunarLander | Policy rules for dNLRLc
---|---
| mean reward:
---
$162.1\pm 110.9$
| doNothing()
---
$:-(CoordY<1.28\wedge LinearVelocX<-0.14)$
$:-([0.80]CoordX>-1.53\wedge[0.76]CoordX<1.70\wedge LinearVelocY<-0.12$
$\wedge[0.66]Angle>-1.60)$
$:-([0.82]CoordX<1.70\wedge LinearVelocY<-0.04\wedge RightLegContactFalse)$
fireLeft())
$:-([0.65]CoordX>-1.52\wedge LinearVelocX<-0.14\wedge[0.73]Angle>-0.09)$
$:-([0.72]CoordX<1.70\wedge[0.68]CoordY>-0.99\wedge LinearVelocX<1.93\wedge
Angle>-0.09)$
fireMain()
$:-(LinearVelocX>0.20\wedge Angle<0.17\wedge AngularVeloc<0.10)$
$:-([0.58]CoordY>-0.91\wedge LinearVelocY>-0.11)$
$:-([0.75]CoordX<1.70\wedge LinearVelocX<-0.11\wedge Angle>0.01\wedge
AngularVeloc>-0.14)$
fireRight()
$:-(LinearVelocX>-0.63\wedge Angle>-1.60\wedge Angle<0.17\wedge
AngularVeloc<-0.02)$
$:-([0.60]CoordX>-1.53\wedge Angle<0.17\wedge[0.94]AngularVeloc>-2.58$
$\wedge LeftLegContactTrue)$
$:-(CoordY<1.28\wedge LinearVelocX>-0.05\wedge Angle<0.32\wedge
AngularVeloc>-1.98$
$\wedge RightLegContactFalse)$
$:-(LinearVelocX>-0.05\wedge Angle>-1.60\wedge Angle<0.17)$
Table 3: The FOL policy rules for the dNLRLc agent, trained using SAC.
Extracted continuous bounded predicates define rules for each action.

In Figure 4, we can observe the results comparing the performance of the
dNLRLnlc agent. The inclusion of the $sine$ is a very minor addition, but one
that does take into account the state transition mechanics. Performance is
better than dNLRLc in later episodes. The moving average shows mean rewards
rising above $100.0$ consistently. The rules in the predicate action policy
for the Lunar Lander environment are given in table 4. For each discrete
action, all rules pass the parameter threshold for the disjunction layer
except for a single rule associated with $fireLeft()$. For the predicate
actions fireLeft() and fireRight(), which have four rules each, we observe the
majority of definitions include the $sine$ transformed state feature $\theta$.

Figure 4: Episode scores for SAC algorithms relying on dNLRLnlc. LunarLander | Policy rules for dNLRLnlc
---|---
| mean reward:
---
$133.4\pm 113.9$
| doNothing()
---
$:-(CoordX<1.66\wedge LinearVelocY<-0.61)$
$:-([0.90]CoordX<1.66\wedge[0.65]CoordY>-0.98\wedge[0.91]CoordY<1.42\wedge
LeftLegContactFalse))$
$:-(LinearVelocY<-0.22)$
$:-([0.61]CoordX<1.51\wedge AngleSine<-0.25)$
$:-([0.86]CoordY<2.12\wedge AngularVeloc>0.23)$
fireLeft())
$:-(LinearVelocX<-0.15)$
$:-([0.88]CoordY>-0.98\wedge
LinearVelocX<0.20\wedge[0.56]LinearVelocY<0.92\wedge AngleSine>-0.08)$
$:-([0.53]AngleSine>-1.00\wedge AngleSine>-0.08\wedge AngularVeloc>0.07)$
$:-[0.65]([0.52]CoordY>-0.98\wedge
CoordY<0.14\wedge[0.72]LinearVelocY>-2.32,[0.68]AngleSine>-1.00$
$\wedge AngleSine>-0.08\wedge[0.85]AngularVeloc<1.88)$
fireMain()
$:-(LinearVelocX<-0.15\wedge
LinearVelocX<-0.58\wedge[0.71]LinearVelocX<0.20\wedge AngleSine>-0.08\wedge$
$AngularVeloc>-1.12\wedge LeftLegContactFalse)$
$:-(CoordX>-1.06\wedge LinearVelocY>-0.11)$
$:-(CoordY>0.34\wedge LinearVelocY>-0.42\wedge[0.62]AngularVeloc>-1.12)$
$:-(LinearVelocX>0.42\wedge AngleSine<-0.25\wedge AngularVeloc>-1.12)$
fireRight()
$:-(LinearVelocX>-0.25\wedge LinearVelocY<-0.61\wedge
AngleSine<0.37,[0.73]AngularVeloc<2.08)$
$:-(CoordX>0.36\wedge
CoordY<0.14\wedge[0.64]LinearVelocY>-2.32\wedge[0.70]AngleSine<0.37$
$\wedge AngleSine<0.25\wedge[0.74]AngularVeloc>-1.12)$
$:-([0.52]CoordY>-0.91\wedge AngleSine<0.25\wedge AngularVeloc<0.05)$
$:-(CoordX>0.36\wedge CoordY<1.42\wedge AngleSine<0.25\wedge
AngularVeloc<1.88\wedge LeftLegContactFalse)$
Table 4: The FOL policy rules for the dNLRLnlc agent, trained using SAC.
Extracted continuous bounded predicates define rules for each action with the
inclusion of non-linear continuous bounded predicates.

####  4.2.3 Discussion

The performance of both the dNLRLc and dNLRLnlc agents is comparable. In the
CartPole problem, see Figure 2, the dNLRLnlc appears more stable in later
episodes than dNLRLc as evidenced by the higher occurrence of total rewards of
$300$. The Lunar Lander problem, with a larger action and state space, proved
more challenging for the dNLRL agents. Both agents were unable to produce mean
rewards of $200$, although dNLRLnlc achieved higher rewards than dNLRLc in
later episodes. While the deviation was high, indicating both dNLRLc and
dNLRLnlc landed successfully on occasion, this also signified that the Lander
would crash periodically. The various hyperparameters were also found to
impact the model. Specifically, the instantiation of the binning scheme would
sway performance. We note, that in [8], equal-width binning was not
necessarily an optimal approach for the Lunar Lander except in specific
regions with respect to the positional axis, and we leave this investigation
for future research. Multiple trials of both agents were conducted, with
performance across all episodes averaging around $100$ for dNLRLc and slightly
higher for dNLRLnlc, and performance fluctuations resulted in a variety of
final policy rules. The position of the lander along the $X$ and $Y$ axis, as
signified by atoms associated with $CoordX$ or $CoordY$, would often have
corresponding conjunction weights, indicating uncertainty in these predicates.
The fluctuating bounded values during training might have caused confusion for
the agent, as the position of the lander is a significant factor in landing
success. To address this issue, future investigations could explore
alternative schemes for training the bound weights, and consider adding
additional non-linear transformation predicates or operation predicates
between state features.

##  5 Conclusion

We aimed to incorporate both continuous and non-linear interpretations of dNL
networks into an RRL framework, creating a dNL-ILP based agent that can learn
in dynamic continuous environments. SAC was found to be the best among RL
algorithms for evaluating the agent. The agent produced policies incorporating
continuous and non-linear continuous predicate functions and was the first to
successfully integrate ILP-based reasoning, RRL, and learning in dynamic
continuous settings. The Lunar Lander problem was more challenging for the
dNLRL agents, resulting in a high deviation from the mean, but our dNLRLc and
dNLRLnlc agents still provide a promising starting point for ILP and RL
research in continuous domains.

## References

  * [1]
  * [2] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang & Wojciech Zaremba (2016): _OpenAI Gym_.  Available at http://arxiv.org/abs/1606.01540.
  * [3] Petros Christodoulou (2019): _Soft Actor-Critic for Discrete Action Settings_ , 10.48550/ARXIV.1910.07207.  Available at https://arxiv.org/abs/1910.07207.
  * [4] Honghua Dong, Jiayuan Mao, Tian Lin, Chong Wang, Lihong Li & Denny Zhou (2019): _Neural Logic Machines_ , 10.48550/arXiv.1904.11694.  arXiv:https://arxiv.org/abs/1904.11694.
  * [5] Kurt Driessens (2010): _Relational Reinforcement Learning_ , 10.1007/978-0-387-30164-8_721.
  * [6] Dheeru Dua & Casey Graff (2017): _UCI Machine Learning Repository_.  Available at http://archive.ics.uci.edu/ml.
  * [7] Richard Evans & Edward Grefenstette (2017): _Learning Explanatory Rules from Noisy Data_ , 10.48550/arXiv.1711.04574.  arXiv:https://arxiv.org/abs/1711.04574.
  * [8] Soham Gadgil, Yunfeng Xin & Chengzhe Xu (2020): _Solving The Lunar Lander Problem under Uncertainty using Reinforcement Learning_ , 10.48550/arXiv.2011.11850.  arXiv:https://arxiv.org/abs/2011.11850.
  * [9] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel & Sergey Levine (2018): _Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor_ , 10.48550/ARXIV.1801.01290.  Available at https://arxiv.org/abs/1801.01290.
  * [10] Rodrigo Toro Icarte, Toryn Klassen, Richard Valenzano & Sheila McIlraith (2018): _Using Reward Machines for High-Level Task Specification and Decomposition in Reinforcement Learning_.  Available at http://proceedings.mlr.press/v80/icarte18a.html.
  * [11] Zhengyao Jiang & Shan Luo (2019): _Neural Logic Reinforcement Learning_ , 10.48550/arXiv.1904.10729.  arXiv:https://arxiv.org/abs/1904.10729.
  * [12] Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini & Shane Legg (2018): _Scalable agent alignment via reward modeling: a research direction_ , 10.48550/arXiv.1811.07871.  arXiv:https://arxiv.org/abs/1811.07871.
  * [13] Nikita I. Lytkin, Risto Miikkulainen & Nate Kohl (2005): _Reinforcement Learning Benchmarks and Bake-offs II A workshop at the 2005 NIPS conference_.
  * [14] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver & Koray Kavukcuoglu (2016): _Asynchronous Methods for Deep Reinforcement Learning_ , 10.48550/ARXIV.1602.01783.  Available at https://arxiv.org/abs/1602.01783.
  * [15] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra & Martin Riedmiller (2013): _Playing Atari with Deep Reinforcement Learning_ , 10.48550/ARXIV.1312.5602.  Available at https://arxiv.org/abs/1312.5602.
  * [16] Stephen Muggleton & Luc de Raedt (1994): _Inductive Logic Programming: Theory and methods_ , 10.1016/0743-1066(94)90035-3.  Available at http://www.sciencedirect.com/science/article/pii/0743106694900353.  Special Issue: Ten Years of Logic Programming.
  * [17] Ali Payani & Faramarz Fekri (2019): _Inductive Logic Programming via Differentiable Deep Neural Logic Networks_ , 10.48550/arXiv.1906.03523.  arXiv:https://arxiv.org/abs/1906.03523.
  * [18] Ali Payani & Faramarz Fekri (2020): _Incorporating Relational Background Knowledge into Reinforcement Learning via Differentiable Inductive Logic Programming_ , 10.48550/arXiv.2003.10386.  arXiv:https://arxiv.org/abs/2003.10386.
  * [19] Melrose Roderick, Christopher Grimm & Stefanie Tellex (2017): _Deep Abstract Q-Networks_ , 10.48550/arXiv.1710.00459.  arXiv:https://arxiv.org/abs/1710.00459.
  * [20] Andrew Silva, Taylor Killian, Ivan Dario Jimenez Rodriguez, Sung-Hyun Son & Matthew Gombolay (2020): _Optimization Methods for Interpretable Differentiable Decision Trees in Reinforcement Learning_ , 10.48550/arXiv.1903.09338.  arXiv:https://arxiv.org/abs/1903.09338.
  * [21] Richard S Sutton, Andrew G Barto et al. (1998): _Introduction to reinforcement learning_.
  * [22] Richard S Sutton, David McAllester, Satinder Singh & Yishay Mansour (1999): _Policy Gradient Methods for Reinforcement Learning with Function Approximation_.  Available at https://proceedings.neurips.cc/paper/1999/file/464d828b85b0bed98e80ade0a5c43b0f-Paper.pdf.
  * [23] Yuan Yang & Le Song (2019): _Learn to Explain Efficiently via Neural Logic Inductive Learning_ , 10.48550/arXiv.1910.02481.  arXiv:https://arxiv.org/abs/1910.02481.
  * [24] Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, Murray Shanahan, Victoria Langston, Razvan Pascanu, Matthew Botvinick, Oriol Vinyals & Peter Battaglia (2018): _Relational Deep Reinforcement Learning_ , 10.48550/arXiv.1806.01830.  arXiv:https://arxiv.org/abs/1806.01830.