\section{Introduction}

Emerging powered prosthetic legs have the potential to restore normative gait biomechanics in individuals with limb loss. The most prevalent control framework for powered prosthetic devices has three levels~\cite{Tucker2015, Varol2009}. The high-level controller is responsible for recognizing the user's intent (e.g., level walking, sitting, stair ascent, stair descent) and estimating the environment's parameters (e.g., ground slope, stair height) to provide input to the mid-level controller. The mid-level controller calculates the reference prosthetic joint angles or torques with information from the high-level controller. Finally, the low-level controller tracks the reference signal and drives the motors so that the prosthesis operates according to the user's intentions and task conditions.

\begin{figure}[t]
    \centering
    \includegraphics[width=.99\columnwidth]{plots/motivation.pdf}

    \caption{An intent prediction model based on a convolutional neural network trained on a large dataset will not perform well when validating new subject data. We study a novel application of transfer learning to improve the performance of the intent prediction model with limited data from a new subject.}
    \label{fig:my_label}
\end{figure}

It is evident that accurately and effectively predicting user intent is crucial for the prosthetic controller to perform well without negatively impacting the user. For instance, an incorrect classification between specific activities, such as level walking and stair descent, could result in the user falling down the stairs. Recent techniques predict user intent continuously by combining supervised learning algorithms with real-time data from onboard mechanical sensors such as an Inertial Measurement Unit (IMU)~\cite{Bhakta2020, Cheng2021, Bruinsma2021, Su2019}. These investigations highlight the potential of predicting user intent by using modern supervised learning algorithms with onboard sensor data. These algorithms can be divided into two categories: feature engineering and feature learning. For the feature engineering method, manually computed statistical values from sensor data are fed into supervised learning models such as Linear Discriminant Analysis (LD) or Support Vector Machines (SVM). On the contrary, the feature learning method will automatically extract features from sensor data using deep neural networks such as convolutional neural network (CNN), recurrent neural network (RNN), or long short-term memory (LSTM)~\cite{Kuangen2019, Jindong2017}.

Compared to studies using data collected from able-bodied individuals, fewer studies on locomotion intent prediction use data collected from individuals with a limb amputation~\cite{Bruinsma2021}. Some studies use data from amputee participants, but the number of subjects is small, making it difficult to generalize and extend intent prediction models to other users~\cite{Bhakta2020}. Different studies use data from able-bodied datasets, which have more subjects~\cite{Kang2022, Lee2020}. However, when evaluating intent prediction models with subjects not included in the trained dataset, the accuracy is significantly lower than when the trained dataset contains test subject data. In clinical environments, it is challenging to collect sufficient data from amputee patients to train a supervised learning model; therefore, a general intent prediction model that can be applied to the novel subject is desired.

This paper proposes a method to improve the accuracy of a base CNN model evaluated on a novel subject by using transfer learning, which requires only a few strides of training data per activity from that new user. The transfer learning method is driven by the fact that humans can learn more quickly, efficiently, and effectively by transferring knowledge from earlier learning experiences. Transfer learning can reduce the quantity of data needed to train a deep CNN model by using techniques such as feature extraction or modifying parameters from a pre-trained model~\cite{Abolfazl2021, Kaur2021, Pan2010}.

We use a public able-bodied dataset~\cite{Hu2018} to train and evaluate the intent prediction model based on a deep CNN. The subject-dependent and subject-independent models are developed and trained using data from various sensor options that can be obtained from a prosthetic leg and amputee user. The subject-independent model is transferred to a new model (referred to as the transfer learning model) by freezing the convolutional layers and training the fully connected layers with a small sample of test subject data. This study demonstrates that the error rate of the subject-dependent model is $1.20\pm 0.002$\%, the subject-independent model is $12.87\pm 0.077$\%, and the transfer learning model using 10\% test subject data is $5.28\pm 0.029$\%. The transfer learning model's accuracy increases proportional to the quantity of data included in the training.

In addition, we demonstrate that the accuracy of using data from a single IMU on the thigh of one leg (e.g., the prosthesis) is comparable to using other sensor configurations. This is significant for practical applications because we do not need to add sensors to the sound leg nor rely on prosthetic leg sensors that are directly affected by the controller rather than the user (i.e., encoders, shank IMU).

The main contributions of this paper are:
\begin{enumerate}

    \item Application of transfer learning to a subject-independent model to increase the accuracy of locomotion intent prediction for new subjects with just a few test subject strides per activity.

    \item Analysis of various sensor configurations suggesting that a thigh IMU on one leg (e.g., the prosthesis) can be used as the input signal to predict locomotion intent with comparable accuracy to other options.
\end{enumerate}

The remaining sections are organized as follows. Section \ref{Method} describes the dataset and intent prediction models. The results are summarized in Section \ref{Results} and discussed in Section \ref{Disscussions}. In Section \ref{Conclusions}, concluding remarks are drawn.

\section{Methods} \label{Method}
This section describes the deep neural network architectures for locomotion intent prediction based on raw data obtained from able-bodied individuals using a variety of sensor options.

\subsection{Dataset}
In this study, we use a publicly available dataset titled ENcyclopedia  of  Able-bodied  Bilateral  Lower  Limb  Locomotor  Signals  (ENABL3S)~\cite{Hu2018}
to train a deep convolutional neural network model for predicting user intentions. The dataset includes various types of sensor signals mounted on ten different able-bodied subjects, such as wearable electrogoniometers (GONIO), surface electromyography (EMG) and inertial measurement unit (IMU) sensors. Since mechanical sensors are commonly available on powered prosthetic legs \cite{Torrealba2019}, our study focuses on using features from IMU and GONIO to predict the intent of users.

The GONIO sensor measures the knee and ankle angle/velocity, which would be measured by joint encoders in a powered prosthetic leg. In the dataset, each subject repeats approximately 25 trials of a two-stage circuit comprising of sitting (S), standing (St), level walking (LW), stair ascent (SA) (stage 1) or ramp ascent (RA) (stage 2), LW, ramp descent (RD) (stage 1) or stair descent (SD) (stage 2), LW, and S. The slope of the ramp is 10 degrees, and the staircase has four steps. This dataset labels the locomotion modes manually with a key fob.
We exclude one of the subjects (AB186) from the dataset because we conjecture there is a problem with the thigh IMU signals, e.g., the sign of the y-axis thigh angular velocity is opposite to that of the other subjects performing the same mode.

\subsection{Data Preparation}
\subsubsection{Data input}

We analyze and evaluate the selection of features gathered from various sensors to predict the user's intentions. Four scenarios are compared, including (i) a thigh IMU on one leg (e.g., the prosthesis), (ii) thigh IMUs on both legs, (iii) onboard sensors available to a prosthesis (i.e., IMUs on the thigh and shank, encoders (ENCs) on the ankle and knee joints), and (iv) all sensors (i.e., including sensors on the prosthesis and an IMU on the healthy leg). The ENC signal correlates with the GONIO sensor signal in the dataset. Each IMU has a 3-axis gyroscope and a 3-axis accelerometer, providing three signals for angular velocity and three for linear acceleration, respectively.

Previous studies relied on a prior gait event signal window as a data sample to be input into the model~\cite{Lee2020, Bhakta2020}. However, this strategy can be difficult to implement in practice due to the difficulty of identifying gait events~\cite{Xu2021}, and missed gait events could cause missed mode transitions. In this study, a 500 ms sliding window similar to~\cite{Bruinsma2021, Kang2022} is sampled for model input, and the label is the locomotion mode in the next time step.

\subsubsection{Data splitting}
We evaluate three models for predicting user intent: subject-dependent, subject-independent, and transfer learning.

a) Subject dependent: We train and test the deep learning models using the data from the same subject. For each subject, we split that subject's data into training (80\%), validation (10\%), and testing (10\%)

b) Subject independent: We use leave-one-subject-out cross-validation to train the deep learning models on the data from a set of subjects and test the model on a new untrained subject. The training and validation set consists of data from eight subjects divided into 80\% for training and 20\% for validation. The remaining subject data is used as a test set. This process was repeated nine times until all subjects have been tested.

c) Transfer learning: The transfer learning method is applied to a subject-independent model that uses a portion of the test subject's data for the training process. In order to validate the effect of data size, we use 5\% (3\% for training and 2\% for validation), 10\% (7\% for training and 3\% for validation), 15\% (10\% for training and 5\% for validation), and 20\% (15\% for training and 5\% for validation) from the test subject data for model training. We also apply this process to all subjects with the corresponding subject-independent model.

\subsubsection{Data output}
In our previous studies~\cite{Best2021, cheng2022}, we merge standing and ramp walking as level walking with variable velocities and slopes, thus we simply need to predict sit (S), level walking (LW), stair ascent (SA), and stair descent (SD), as well as the transitions between them: S2W, W2S, W2SA, SA2W, SD2W. In total we need to classify ten different locomotion modes.

\subsection{Deep Neural Networks Model Architectures}
\subsubsection{Base convolutional neural network}

We build the intent prediction models with a deep convolutional neural network. We train and validate the CNN models with two different methods, i.e., subject-dependent and subject-independent. CNN model is developed with the VGG network method~\cite{Simonyan2014, zhang2021dive}. The model is built from VGG blocks, with each VGG block beginning with a convolutional layer and ending with a max pooling layer as shown in Fig.~\ref{fig:vgg_block}. The convolutional layer is designed to extract complex features from the input data, and the one-dimensional max pooling layer is added after each block to reduce the data sample size. A one-dimensional kernel is used to enhance computational performance in terms of memory usage and execution time. The rectified linear activation function (ReLU) is used for the convolutional layer, together with a dropout mechanism to prevent overfitting.

\begin{figure}
  \centering
  \includegraphics[width=0.8\columnwidth]{plots/vgg_block.pdf}
  \caption{\label{fig:vgg_block} VGG block architectures consist of one convolutional (Conv) layer with 3×1 kernels with padding of 1 and stride of 1 (keeping the height and width) followed by ReLU activation function, dropout of 0.2, and a 2×1 max-pooling layer with a stride of 2 (halving height after each block). The convolutional part of the network connects several VGG blocks.
  The input channels to the first block correspond to sensor signals. For example, each IMU sensor has six channels and each encoder has two channels.}
\end{figure}

\begin{figure}
  \centering
  \includegraphics[width=0.9\columnwidth]{plots/dep_cnn.pdf}
  \caption{\label{fig:dep_cnn} Subject-dependent model architectures include four continuously connected VGG blocks and three fully connected layers. The output channels for each VGG block are 64, 128, 246, and 512. The first fully connected layer flatten the output from convolution layers. The output size of second fully connected layer is 1024 and the last layer is ten, which is the number of labeled modes.}
\end{figure}

As depicted in Fig. \ref{fig:dep_cnn}, the model employed for the subject-dependent method consists of four continuously connected VGG blocks. The number of input channels in the first block depends on the number of sensors used; for instance, an IMU sensor has six channels, whereas an encoder has two. Each VGG block's output channels correspond to the input of the following block. A series of 64–128–256–512 output channels are developed for the subject-dependent model. As shown in Figure \ref{fig:ind_cnn}, the subject-independent model uses one more VGG block than the subject-dependent model, with output channels of 1028. The VGG blocks are interconnected to construct a system of convolutional layers responsible for extracting data-descriptive features.

Features learned by the convolutional layers are then fed forward into the fully connected (FC) layers, which are used to map those features into the pre-labeled ten locomotion modes. Similar to convolutional layers, we also involve dropout methods and the ReLU activation function in the first two fully connected layers to prevent overfitting. The second fully connected layer in the subject-dependent model has 1024 outputs while subject-independent has 2048 outputs. Both models have the fully connected output layer with 10 outputs, which is the number of labeled locomotion modes.

\subsubsection{Transfer learning}

The subject-independent model is then refined via a transfer learning approach by freezing the learned parameters in the convolutional layers and relearning the parameters in the fully connected layers.

\begin{figure}
  \centering
  \includegraphics[width=0.9\columnwidth]{plots/ind_cnn.pdf}
  \caption{\label{fig:ind_cnn} Subject-independent model architectures included five continuously connected VGG blocks and three fully connected layers. The output channels for each VGG block are 64, 128, 246, 512, and 1024. The first fully connected layer flatten the output from convolution layers. The output size of second fully connected layer is 1024 and the last layer is ten, which is number of labeled modes.}
\end{figure}

\subsection{Deep Neural Networks Model Training}
This section covers how to train deep convolutional neural networks. The training was conducted on a PC with 2 NVIDIA TITAN RTX GPUs. The model was constructed and trained using the Pytorch Lightning Library.

\subsubsection{Loss function}

We set up a classification problem by minimizing the cross entropy loss defined as $$L=-\sum_{c=1}^C w_c \log \frac{\exp \left(x_c\right)}{\sum_{i=1}^C \exp \left(x_i\right)} y_c,$$ where \(y_c\) is the target, \(x_c\) is the input, \(w_c\) is the class weight, and \(C\) is the number of classes.

The number of data samples from different classes (i.e., locomotion modes) is unbalanced in the ENABL3S dataset. To address this issue, we assign different weights to each class based on the percentage of strides for each activity and account that into the loss function.

\subsubsection{Hyperparameter}

We investigate various CNN model hyperparameters for optimal performance. The number of epochs, batch size, and learning rate are involved parameters in model training. Table \ref{training_params} contains the parameters chosen for high accuracy.

The training data are divided into data sample batches for model training. The batches are used in succession until all training data is consumed to conclude an epoch. The process is repeated numerous times, and the greater the number of epochs, the higher the training accuracy but also the longer the training time.

Learning rate is another parameter for model training. The higher the learning rate, the quicker the training, but if it is too high, the model cannot converge. Conversely, if the learning rate is too low, the loss function may decrease to a local minimum, or the training time may become too long.

\begin{table}[t]
\caption{Optimal hyperparameter (i.e., number of epoch, batch size, and learning rate) for different CNN models included subject-dependent (Dep.), subject-independent (i.e., Ind.), and transfer learning (i.e., Transfer).}
\label{training_params}
\footnotesize
\begin{center}
\begin{tabular}{c|c|c|c}
\hline
& Dep. & Ind. & Transfer \\
\hline
Epoch & 30& 35& 100\\
\hline
Batch size& 512& 1024& 256\\
\hline
Learning rate & 0.0001& 0.00015& 0.0001\\
\hline
\end{tabular}
\end{center}
\end{table}

\subsubsection{Optimizer}
Depending on the optimizer's formula, the optimizer can be interpreted as a mathematical function that adjusts the neural network's weights using gradients and additional information. We use the Adaptive Moment Estimation (Adam)~\cite{Kingma2014} optimizer for the subject-dependent and subject-independent models and the stochastic gradient descent (SGD) algorithm for the transfer learning model in this study. Adam is the default optimization algorithm for training deep learning models in numerous studies due to its faster training speed. However, when the transfer learning method is used, the amount of training data and number of parameters are small, so the more efficient SGD method is the most suitable option~\cite{Sun2019, Sun2019_2}.

\subsection{Evaluation}

Similar to previous studies on locomotion intent classifiers~\cite{Hu2018, Bhakta2020, Cheng2021, Lee2020}, we evaluate the model based on the error rate criterion determined by the proportion of incorrect predictions relative to the total number of test data samples. Moreover, the error rates can be categorized as steady state (SS) and transition (TS). A TS occurs when the current mode differs from the previous mode, whereas a SS occurs when the current mode is identical to the previous mode. The formula for calculating these criteria are:

\begin{equation}
\text{Overall\_error}=1 - \frac{\text{Overall\_correct}}{\text{Overall\_total}},
\end{equation}

\begin{equation}
\text{SS\_error}=1 - \frac{\text{SS\_correct}}{\text{SS\_total}},
\end{equation}

\begin{equation}
\text{TS\_error}=1 - \frac{\text{TS\_correct}}{\text{TS\_total}}.
\end{equation}

In addition, we perform statistical analysis to compare the model performance over two factors: learning condition (DEP, IND, and Transfer Learning) and sensor setup (unilateral thigh, bilateral thigh, prosthetic sensors, and all). To that end, we conduct a two-way repeated measures analysis of variance (ANOVA) with classification error as the dependent variable, and learning condition and sensor setup as independent variables. Then, we followed up with a Bonferroni post-hoc analysis for the significant factors to determine the statistical significant differences (p < 0.05) between each level within each factor. Finally, we ran multiple pairwise comparisons to calculate the significant difference in error rate between each specific pair of learning conditions within each sensor setup (e.g., between unilateral thigh under DEP and unilateral thigh under IND).

\section{Results} \label{Results}
The results of the three proposed models with different sensor configurations are shown in Fig. \ref{fig:mainResults}, where 10\% of the test subject data is used for the transfer learning method. Moreover, Fig. \ref{fig:TransferResults} shows the effect of the amount of test subject data used in the transfer learning method (5\%, 10\%, 15\%, and 20\%).

\begin{figure}[t]
  \centering
  \vspace{-80pt}
  \includegraphics[width=0.48\textwidth]{plots/Final_ICRA.pdf}
  \vspace{-585pt}
  \caption{Three different models (subject-dependent, subject-independent, and transfer learning) are compared in different sensor setups and error types. The unilateral thigh means using one thigh IMU on the prosthesis leg, the bilateral thigh means using both thigh IMU on both legs, prosthetic sensors means using all sensor on the prostheses, and All means using all sensors included thigh IMU on healthy leg and sensors on prosthetic leg. Error types are calculated as in equations (1)-(3). Error bars represent $\pm$ standard error of the mean. Asterisks indicate statistical significance (p < 0.05).}
  \label{fig:mainResults}
\end{figure}

\subsection{Model Comparison}
Table \ref{tbl:result_num} compares the error rates of the intent prediction models with different sensor selections. The overall error rate of the subject-independent model is statistically higher than the subject-dependent model ($p<0.05$), however, the subject-dependent model is not significantly different from the transfer learning model ($p>0.05$). This trend tends to hold with steady state and transition errors for all sensor setups. In addition, transfer learning model demonstrates significantly lower error rate than the subject-independent model only in unilateral thigh and bilateral thigh sensor setups, i.e., no significant differences in Prosthetic and All sensors setups

\begin{table}[t]
\caption{Error rates of the classifiers on different sensor setups and models including subject-independent (Ind.), subject-dependent (Dep.), and transfer learning (Transfer). All results are represented as mean [standard deviation].}
\label{tbl:result_num}
\footnotesize
\begin{center}
\begin{tabular}{c|c|c|c|c}
\hline
Sensor setup & Error & Dep. & Ind. & Transfer. \\
\hline
Unilateral Thigh & Overall & 1.20[.002] & 12.87[.077] & 5.28[.029] \\
 & SS & 0.38[.001] & 9.64[.062] & 4.03[.023] \\
 & TS & 5.01[.011] & 29.82[.175] & 11.94[.060] \\
 \hline
Bilateral Thigh & Overall & 0.74[.002] & 11.59[.076] & 3.57[.020] \\
 & SS & 0.18[.001] & 8.72[.061] & 2.57[.016] \\
 & TS & 3.46[.010] & 26.06[.166] & 8.74[.045] \\
 \hline
Prostheses Leg & Overall & 1.06[.003] & 9.19[.053] & 3.71[.019] \\
 & SS & 0.23[.001] & 6.03[.038] & 2.43[.014] \\
 & TS & 5.00[.017] & 25.04[.138] & 10.24[.045] \\
 \hline
All & Overall & 0.84[.002] & 8.41[.049] & 3.45[.018] \\
 & SS & 0.21[.001] & 5.65[.038] & 2.36[.014] \\
 & TS & 3.84[.012] & 22.52[.128] & 9.11[.045] \\
 \hline
\end{tabular}

\end{center}
\end{table}

\begin{figure}
  \centering
  \includegraphics[width=1\columnwidth]{plots/transfer.JPG}
  \caption{\label{fig:TransferResults} Prediction error for transfer learning at different percentages of trained dataset from test subject.}
\end{figure}

\subsection{Sensor Selections Comparison}

When we compare models with different sensor choices, ANOVA found no statistical difference between them. However, it is still possible to see some differences between these models. For the subject-independent model, the more sensors used, the lower the overall error rate (Unilateral Thigh: 12.87\%, Bilateral Thigh: 11.59\%, Prostheses Leg: 9.19\%, All: 8.41\%). The same goes for steady-state and transitional error rates. However, with the subject-dependent model, choosing the bilateral thigh sensor option gives the lowest error rate (Overall: 0.74\%, Steady-state: 0.18\%, Transition: 3.46\%). For the transfer learning model, choosing to use all sensors gives the lowest overall and steady-state error rates (Overall: 3.45\%, Steady-state: 2.35\%). But the transfer learning model with Bilateral Thigh sensor selection gives the lowest transition error rate (8.74\%).

\subsection{Data for Transfer Learning Comparison}

\begin{table}[t]
\caption{Estimation of number of strides for 10\% of a subject data. These values are approximated by 10\% of the averaged total number of strides for ten subjects reported in~\cite{Hu2018}.}
\label{tbl:10percentsub}
\footnotesize
\begin{center}
\begin{tabular}{c|c|c}
\hline
 &  & Number of strides \\
 \hline
Steady-state & Level walking (LW) & 45 \\
 & Ramp ascent (RA) & 14 \\
 & Ramp descent (RD) & 18 \\
 & Stair ascent (SA) & 5 \\
 & Stair descent (SD) & 5 \\
 \hline
Transition & LW to RA & 2 \\
 & LW to RD & 2 \\
 & LW to SA & 2 \\
 & LW to SD & 2 \\
 & RA to LW & 2 \\
 & RD to LW & 2 \\
 & SA to LW & 2 \\
 & SD to LW & 2 \\
 \hline
Total &  & 106 \\
\hline
\end{tabular}
\end{center}
\end{table}

Fig. \ref{fig:TransferResults} shows the accuracy of user intent prediction with different percentages of data used in transfer learning for the test subject. Generally, the prediction errors decreased with more training data available from the test subject, and this trend is consistent across different sensor setups.
Since the ENABL3s dataset contains 25 repetitions of a circuit for each subject, 20\% of the test subject data corresponds to approximately five repetitions and 5\% to more than one repetition. More specifically, Table \ref{tbl:10percentsub} shows an estimation of the number of strides for each locomotion mode with 10\% of one subject's data from~\cite{Hu2018}. This means that 10\% of one subject's data is about 106 strides, including two strides for each mode transition.

\section{Discussions} \label{Disscussions}
We propose using transfer learning to enhance a subject-independent model's ability to predict the locomotion intent of a new subject. In addition, we examine various sensor options for predicting user intent to identify a practical solution.
\subsection{Comparison to Prior Studies}
In this study, we explore the novel application of transfer learning to studies of locomotion intent prediction using mechanical sensors.
It is difficult to directly compare our intent prediction models to previous research due to differences in the training dataset, labeled modes, sensor setups, and evaluation methods. However, we can compare our subject-dependent and subject-independent models to other studies to give context for our transfer learning results. Our overall accuracy in predicting locomotion intent using a single thigh IMU (DEP error rate: 1.2\%) is comparable to the study using the WaveNet model with the same sensor input (DEP error rate: 2.12\%)~\cite {Lu2020}. In addition, our subject-dependent error rate using two thigh IMUs (DEP error rate: 0.74\%) outperforms the study using the same dataset (DEP error rate using only IMU sensors: 2.58\%)~\cite{Lee2020}. Our overall error rate with prosthetic sensor setup (DEP: 1.06\%, IND: 9.19\%) are comparable to the study using XBoost for data from amputees with all sensors on the prosthesis (DEP error: 2.93 \%, IND error: 10,12 \%)~\cite{Bhakta2020}. The similarity of the DEP and IND models compared with previous studies demonstrate that our transfer learning results are meaningful.

\subsection{Implementation of Transfer Learning}
This study demonstrates the effectiveness of transfer learning in improving the accuracy of intent prediction for novel subjects. With the abundance of data collected from able-bodied individuals~\cite{Camargo2021, Reznick2021}, it is straightforward to construct a generalized intent prediction model. In studies on image classification, numerous models are built with vast quantities of data and are made available for use in other research. Then, other researchers who wish to apply image classification to a significantly smaller dataset can use transfer learning for the public pre-trained models. We propose a similar approach to building an intent prediction model by making the dataset public so that various deep learning models can be built and compared. The pre-trained models can then be used based on the actual application, such as constructing an intent prediction model for prosthesis controllers.
In a clinical environment, prosthesis controllers can be tuned to individuals in a pre-designed environment with assistance from a prosthetist~\cite{quintero2018intuitive}. The application of transfer learning can enhance the intent prediction model for individuals operating in an unknown environment. This study demonstrates that the prediction model has improved significantly with about 100 strides, of which only about two are for each transition task (the most difficult tasks to perform).

\subsection{Sensor Selection}
Previous research on predicting locomotion intent for prosthetic control has used all the sensors available on the device~\cite{Bhakta2020}. The prediction accuracy in the simulation is high due to the use of numerous features from sensors. However, in practice, sensors attached to the prosthesis (such as encoders, shank IMU, and foot IMU) are directly influenced by the controller, so they may send signals that differ from the user's intent if the controller is malfunctioning. Furthermore, different research groups use different control strategies for the prosthesis to control the knee and ankle joints, resulting in different extracted features from the sensors on the prosthesis. Therefore, it is challenging to use a model trained on data collected by one research group for a controller developed by another and signal sources directly controlled by the amputee's body are best able to inform user intent. For example. the IMU mounted at the connection point between the prosthetic knee and the amputee socket (i.e., the thigh IMU) can be used for this purpose~\cite{Bartlett2018, Chinimilli2019, Cheng2021}.
However, adding additional sensors (e.g., another thigh IMU) to the user's sound leg is a potential barrier to clinical deployment. According to the statistical analysis, there is no significant difference between using one thigh IMU mounted on the prosthesis (unilateral thigh) versus using thigh IMUs on both legs (bilateral thigh), hence, it is sufficient to predict user intent accurately by using the on-board thigh IMU on the prosthesis.

\subsection{Limitations}
Our study is limited by the small number of subjects in the dataset (N=9), resulting in a high error rate for the subject-independent model, which affects the accuracy of the transfer learning model. The accuracy of the models can be enhanced as more comprehensive datasets become available.
Similarly, this study only uses data from able-bodied subjects, therefore, the ability for the model to adapt to new amputee subjects requires further validation. In the future work, we will use additional data from amputee subjects to increase the accuracy of the subject-independent model, allowing the transfer learning model to predict amputees intentions with greater accuracy.

\section{Conclusions} \label{Conclusions}
Our study investigated transfer learning to improve the accuracy of a CNN-based model for predicting the locomotion intent of new subjects. We found that the transfer learning method significantly decreased the error rate of the user-independent model with only about 100 strides (with two strides for each transition task) from the test subject. This error rate decreases as more data from the test subject is used to train the transferred model. We also tested the intent prediction model with different sensor configurations. We concluded that using a thigh IMU on a single leg (e.g., the prosthesis) will provide accuracy comparable to other sensor setups.

{\small
\balance
\bibliographystyle{IEEEtran}


\section{Experiments}
\subsection{Datasets and Evaluation Metrics}\label{usod_data}
\subsubsection{\textbf{Training and evaluation data}} We use the USOD10K dataset~\cite{USOD10K} in our experiments; it contains RGB images and ground truth depth maps for various underwater scenes captured at a pixel resolution of $640\times480$. The dataset contains $9229$ training samples and $1026$ testing samples. We also use benchmark images from the Sea-Thru dataset~\cite{akkaynak2019sea} for performance evaluation. Moreover, we test UDepth model on unseen field data collected during oceanic explorations and human-robot collaborative experiments.

\subsubsection{\textbf{Evaluation metrics}} We use four standard metrics~\cite{eigen2014depth} to compare our method against other SOTA models. These error  metrics are defined as follows:
\begin{itemize}
    \item Mean absolute relative error (Abs Rel): $ \frac{1}{n} \sum_p^n \frac{|d_p - \hat{d}_p|}{d}$,
    \item Squared relative error (Sq Rel): $ \frac{1}{n} \sum_d^n \frac{||d_p - \hat{d}_p||^2}{d}$,
    \item Root mean squared error (RMSE): $\sqrt{\frac{1}{n}\sum_d^n(d_p - \hat{d}_p)^2}$, \item $\log_{10}$ error: $\frac{1}{n} \sum_d^n |\log_{10}{(d_p)}-\log_{10}{(\hat{d}_p)}|$,

\end{itemize}
Where $d_p$ is a pixel in depth image $d$, $\hat{d_p}$ is a pixel in the predicted depth image $\hat{d}$, and $n$ is the total number of pixels in $d$ (as well as $\hat{d}$).

\subsection{Implementation Details}
UDepth training is supervised by RGB-D image pairs of natural underwater scenes; the RGB color images are pre-processed to prepare R\textbf{M}I input images and the respective D channel depth maps serve as ground truth. A total of $7383$ images in the USOD10K training dataset are used for training and the remaining $1846$ images are used for validation. We use Pytorch libraries~\cite{paszke2019pytorch} to implement its learning pipeline; AdamW~\cite{loshchilov2017decoupled} is used as the optimizer with an initial learning rate of $0.0001$ and exponential decay adjustment with a multiplicative factor of $0.9$. With a batch size of $4$, Udepth training takes about $5$ minutes per epoch on a single node with NVIDIA\texttrademark~RTX $3080$ GPU. It has about $15.5$M parameters in total: $3.5$M for the MobileNetV2 encoder, $10.9$M for the decoder, and $1.1$M for the mViT-based refiner and the convolutional regressor combined.

\subsection{Qualitative and Quantitative Evaluation}
For baseline performance comparison, we consider the following five models that are widely used for supervised learning of monocular depth estimation: VGG-UNet~\cite{simonyan2014very, ronneberger2015u}, ResNet-UNet~\cite{he2016deep, ronneberger2015u}, Adabins~\cite{bhat2021adabins}, DesneDepth~\cite{alhashim2018high},  and MobileNetV2~\cite{sandler2018mobilenetv2}. We train these models by following their recommended settings on the same pipeline and \textit{train-validation} data splits as UDepth. The qualitative performances of all these models for some samples are illustrated in Fig.~\ref{fig:qual_comp}, while the quantitative results are listed in Table~\ref{tab_usod}-\ref{tab_seathru}.

The two major findings of our experimental analyses are as follows: ($i$) All models exhibit consistently better results across almost all metrics when trained on the R\textbf{M}I space instead of raw RGB inputs, which validates our contribution to domain-aware input space adaptation. ($ii$) The proposed UDepth model outperforms VGG-UNet, ResNet-UNet, DesneDepth, and MobileNetV2 regardless of using R\textbf{M}I or RGB input space. On some metrics, the SOTA model Adabins achieves better scores and generate more accurate visual results. However, Adabins is a significantly heavier model with $78$M+ parameters compared to UDepth, which has only $15.5$M parameters. Hence, our design choices enable UDepth to achieve comparable and often better depth estimation performance than Adabins, despite being over $5\times$ more efficient (at only $20\%$ computational cost).

\begin{table}[t]
\centering
\caption{Quantitative comparison for underwater scene depth estimation performance by UDepth and other SOTA models on USOD10K test set~\cite{USOD10K}. Here, lower scores represent better performance for all metrics in consideration.}
\footnotesize
\begin{tabular}{l||l|l|l|l}
  \Xhline{2\arrayrulewidth}
   \textbf{Model} & Abs Rel & Sq Rel & RMSE & $\log_{10}$ \\

    \Xhline{2\arrayrulewidth}
    VGG-UNet (RGB) & $0.939$ & $0.183$ & $0.150$ & $0.199$ \\

    VGG-UNet (R\textbf{M}I) & $0.886$ & $0.180$ & $0.150$ & $0.197$  \\

    ResNet50-UNet (RGB) & $0.875$ & $0.171$ & $0.152$ & $0.200$  \\

    ResNet50-UNet (R\textbf{M}I) & $0.847$ & $0.164$ & $0.154$ & $0.204$ \\

    AdaBins (RGB) & {$0.617$} & {$0.097$} & {$0.109$} & {$0.154$} \\

    AdaBins (R\textbf{M}I) & {{$0.613$}} & {{$0.090$}} & {{$0.108$}} & {{$0.153$}} \\

    DenseDepth (RGB) & $1.178$ & $0.254$ & $0.168$ & $0.219$ \\

    DenseDepth (R\textbf{M}I) & $1.128$ & $0.241$ & $0.170$ & $0.220$ \\

    MobileNetV2 (RGB) & $0.858$ & $0.156$ & $0.134$ & $0.183$ \\

    MobileNetV2 (R\textbf{M}I) & $0.790$ & $0.144$ & $0.136$ & $0.184$ \\

    UDepth (RGB) & $0.809$ & $0.147$ & {$0.129$} & {$0.176$}  \\

    UDepth (R\textbf{M}I) & {$0.681$} & {$0.123$} & $0.143$ & $0.188$ \\
    \Xhline{2\arrayrulewidth}
\end{tabular}
\label{tab_usod}
\vspace{-2mm}
\end{table}

\begin{table}[t]
\centering
\caption{Quantitative performance comparison on Sea-thru (D3) dataset~\cite{akkaynak2019sea} (same models and metrics in consideration as Table~\ref{tab_usod}).}
\footnotesize
\begin{tabular}{l||l|l|l|l}
  \Xhline{2\arrayrulewidth}
    \textbf{Model} & Abs Rel & Sq Rel & RMSE & $\log_{10}$  \\

    \Xhline{2\arrayrulewidth}
    VGG-UNet (RGB)  & $1.402$ & $0.680$ & $0.439$ & $0.359$  \\

    VGG-UNet (R\textbf{M}I) & $1.271$ & $0.563$ & $0.399$ & $0.337$ \\

    ResNet50-UNet (RGB)  & $1.272$ & $0.563$ & $0.397$ & $0.335$ \\

    ResNet50-UNet (R\textbf{M}I)  & $1.206$ & $0.613$ & $0.420$ & $0.352$  \\

    AdaBins (RGB)  & $1.314$ & $0.663$ & $0.441$ & $0.356$  \\

    AdaBins (R\textbf{M}I)  & $1.258$ & $0.617$ & $0.426$ & $0.346$  \\

    DenseDepth (RGB)  & {$1.073$} & {{$0.448$}} & {{$0.366$}} & {{$0.312$}} \\

    DenseDepth (R\textbf{M}I) & {$1.092$} & {$0.470$} & {$0.370$} & {{$0.312$}} \\

    MobileNetV2 (RGB)  & $1.196$ & $0.565$ & $0.408$ & $0.337$ \\

    MobileNetV2 (R\textbf{M}I)  & $1.170$ & $0.547$ & $0.400$ & $0.329$ \\

    UDepth (RGB)  & $1.304$ & $0.653$ & $0.440$ & $0.355$ \\

    UDepth (R\textbf{M}I)  & {$1.153$} & {$0.514$} & {$0.388$} & {$0.321$} \\
    \Xhline{2\arrayrulewidth}
\end{tabular}
\label{tab_seathru}
\vspace{-1mm}
\end{table}

\begin{figure}[h]
\centering
\includegraphics[width=0.96\linewidth]{imgs/qual_comp_legacy.pdf}
\vspace{-3mm}
\caption{A few qualitative comparisons for generalization performance of UDepth with its two best competitors: Adabins and DenseDepth are shown. These underwater test images are randomly selected from various benchmark datasets~\cite{islam2022svam,li2019underwater}.}
\label{fig:qual_comp_legacy}
\end{figure}

The qualitative results corroborate our analyses; as the visual comparisons of Fig.~\ref{fig:qual_comp} illustrate, the output of UDepth and Adabins are more accurate and consistent with the underwater scene geometry. Udepth does a particularly better job at removing background regions and predicting foreground layers up to scale. Moreover, UDepth demonstrates much better generalization performance compared to other SOTA models, as evident from the results in Table~\ref{tab_seathru}. Here, we use the D3 (reef) scenes from the Sea-Thru dataset~\cite{akkaynak2019sea} for testing only. While Adabins offers good results, UDepth achieves significantly better results across all metrics on unseen test cases. In comparison, DenseDepth model shows slightly better performance on unseen underwater images; however, with $42.6$M parameters, it is about $3\times$ computationally heavier. UDepth's superior performance, particularly for background segmentation and depth continuity on unseen natural underwater scenes are illustrated in Fig.~\ref{fig:qual_comp_legacy}.

\begin{table}
\centering
\caption{Comparison of model parameters, model sizes, and inference rates for three best performing models: on a CPU (Intel\texttrademark~Core i9-3.50GHz) and a GPU (NVIDIA\texttrademark~RTX $3080$).}
\footnotesize
\begin{tabular}{l|r||r|r|r}
  \Xhline{2\arrayrulewidth}
    Model & \# Params & Memory & FPS (CPU) & FPS (GPU) \\ \Xhline{2\arrayrulewidth}
    Adabins & $78.0$ M & $313.8$ MB & $1.98$ & $33.70$ \\ \hline
    DenseDepth & $42.6$ M & $178.3$ MB & $3.06$ & $52.34$ \\ \hline
    UDepth & {$15.6$ M}  & {$62.7$ MB}  & {$13.50$} & {$66.23$} \\ \hline
    \Xhline{2\arrayrulewidth}
\end{tabular}
\label{tab_comp}
\end{table}

\begin{figure}[t]
\centering
\includegraphics[width=0.98\linewidth]{imgs/Domain_Proj.pdf}
\vspace{-1mm}
\caption{A few demonstrations of fast coarse depth prediction by our domain projection step are shown in comparison with the final UDepth estimations. The domain projection facilitates an abstract yet reasonably accurate depth prior, which can be used for fast decision-making by underwater robots. }
\label{fig:domain_proj}
\vspace{-10mm}
\end{figure}

\subsection{Coarse Depth Estimation by Domain Projection}
We further compare the computational efficiency of UDepth with Adabins and DenseDepth in Table~\ref{tab_comp}. UDepth is $3$-$5$ times memory efficient and offers $4.4$-$6.8$ times faster inference rates, with over $66$ FPS inference on a single RTX-3080 GPU and over $13.5$ FPS on CPUs. More importantly, we can extend our domain projection step with guided filtering~\cite{he2012guided} for fast \textit{coarse depth prediction} on low-power embedded devices. We performed thorough evaluations on field experimental data, which suggest that these abstract predictions are reasonable approximations of natural underwater scene depths. As shown in Fig.~\ref{fig:domain_proj}, the filtered domain projections embed useful 3D information about the scenes to facilitate high-level decision-making by visually-guided underwater robots. The end-to-end domain projection and filtering module generates depth maps at $51.5$ FPS rate on NVIDIA\texttrademark~Jetson TX2s and on $7.92$ FPS rate on Raspberry Pi-4s.

\section{Introduction}

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{imgs/fig1.pdf}
\vspace{-5mm}
\caption{The proposed UDepth learning pipeline uses the wavelength-dependent attenuation constraints of underwater domain for an intelligent input-space adaptation from raw RGB images. A computationally light domain projection step then approximates a coarse depth map. The intermediate projections and other pixel-level loss functions guide the end-to-end training of UDepth for fine-grained depth estimations.}
\label{fig:intro}
\vspace{-3mm}
\end{figure}

Underwater depth estimation is the foundation for numerous marine robotics tasks such as autonomous mapping and 3D reconstruction~\cite{digumarti2016underwater}, visual filtering~\cite{islam2020fast}, tracking and servoing~\cite{9340919,raanan2018detection}, navigation~\cite{xanthidis2020navigation,rahman2019svin2}, photometry and imaging technologies~\cite{akkaynak2018revised}, and more. Unlike terrestrial robots, visually guided underwater robots have very few low-cost solutions for dense 3D visual sensing because of the high cost and domain-specific operational complexities involved in deploying underwater LiDARs~\cite{zhou2021overview,mcleod2013autonomous}, RGB-D cameras~\cite{7891971}, or laser scanners~\cite{palomer2019underwater}. Fast monocular depth estimation thus plays a critical role in enabling real-time 3D robot perception and state estimation for important applications such as subsea monitoring and inspection~\cite{shkurti2012multi}, autonomous exploration~\cite{girdhar2014autonomous}, and companion robotics~\cite{islam2021robot}.

Traditional approaches for underwater depth estimation either use active photon counting on Time-of-flight (ToF) cameras~\cite{maccarone2015underwater} or geometric model adaptation on structure light~\cite{massot2015optical,8559463}. Standard depth cameras~\cite{7891971} with water-housing and echo-sounders~\cite{9340919} are also common to acquire underwater scene depths up-to-scale. However, accurate sensory integration and backscatter filtering are major challenges in their field deployments. Another class of approaches holistically evaluates the physical properties of the optics for depth estimation; \ie, they approximate the medium transmission rate, background light, and attenuation coefficient following an underwater image formation model (IFM)~\cite{8084744,7351749,raihan2020depth,chang2018single}, and subsequently estimate depth for each pixel. These methods work reasonably well on clear/non-turbid water-bodies (\eg, Ocean Type I-II)~\cite{song2018rapid}, and are particularly suitable for offline image processing tasks. However, their generalized end-to-end implementations are rather difficult due to the optical parameters' dependency on wavelength and water-body properties~\cite{akkaynak2018revised}.

In recent years, deep visual learning-based methods are used to address these problems with remarkable success~\cite{hambarde2021uw,ye2019deep}. These data-driven methods heavily rely on large-scale datasets, hence researchers often synthetically generate training images~\cite{zhao2021synthesis,zhao2021unpaired,cui2021underwater} due to the lack of RGB-D ground truths for underwater scenes. However, such training pipelines cannot fully capture the complex underwater scene geometry and image statistics, restricting their performance and applicability for practical applications. Additionally, their learning objectives are generally adopted from standard terrestrial depth estimation pipelines or coupled with in-air RGB-D image pairs~\cite{ye2019deep,gupta2019unsupervised}. Consequently, underwater domain knowledge is not incorporated into their learning processes, which leads to poor generalization performance.

In this paper, we formulate a robust and efficient end-to-end model named \textbf{UDepth}, for fast monocular depth estimation by incorporating underwater domain knowledge into its supervised learning pipeline. We adapt a new input space R\textbf{M}I for underwater domain-aware learning, which consists of red channel (R), maximum of green and blue channels (\textbf{M}), and grayscale intensities (I). Due to the wavelength-dependent attenuation constraints of underwater light propagation, the relative differences between R and \textbf{M} channels embed useful scene depth information while the I channel preserves structural contents of the image. Based on R\textbf{M}I, we design a computationally light projection step for pixel-wise \textbf{coarse depth prediction} by a least-squared formulation. We further devise a \textbf{domain projection loss} to enforce the pixel-wise underwater attenuation constraints in the holistic learning process as well.

The proposed UDepth model architecture consists of the highly efficient MobilenetV2 backbone, a lightweight Vision Transformer (mViT)-based global attention module, and a convolutional regressor for fine-grained depth estimation. The end-to-end inference graph of UDepth has only $15.6$M parameters, which is about $80\%$ less than Adabins~\cite{bhat2021adabins} and $66\%$ less than DenseDepth~\cite{alhashim2018high} - two of its closest competitor baselines. As a result, UDepth offers significantly faster inference rates: over $66.7$ FPS on a NVIDIA\texttrademark~RTX $3080$ and $13.3$ FPS on an Intel\texttrademark~Core i9-3.50GHz CPU core. With comprehensive domain-aware learning on $9223$ RGB-D pairs of natural underwater schenes from USOD10K datset~\cite{USOD10K}, UDepth achieves SOTA performance on standard benchmarks despite having such a light architecture. Moreover, we demonstrate that UDepth offers better generalization performance on arbitrary test cases from Sea-Thru dataset~\cite{akkaynak2019sea} and real-world field experimental data.

Furthermore, we demonstrate that our domain projection module with subsequent filtering can be used for coarse depth prediction on low-power embedded devices. Specifically, it runs at $51.5$ FPS on NVIDIA\texttrademark~Jetson TX2s and $7.92$ FPS on Raspberry Pi-4s. More importantly, visual results on field experimental data suggest that its coarse predictions are reasonably accurate and can be used for on-board 3D perception by low-cost underwater robots.

\section{Background \& Related Work}
\subsection{Monocular Depth Estimation Literature}
Traditional methods for monocular depth estimation~\cite{godard2019digging,godard2017unsupervised} focus on graphical models with hand-crafted geometric priors based on the concepts of structure-from-motion, stereo vision, and multi-view feature matching. These classical methods require multi-view correspondences and significant computational power, yet only generate sparse depth information. Such difficulties have paved the way for powerful deep visual learning-based methods~\cite{zhao2020monocular}, which can learn to infer dense depth maps from single RGB images in an end-to-end manner. The state-of-the-art (SOTA) methods of the modern era for monocular depth estimation can be classified as supervised, semi-supervised, and unsupervised approaches.

\subsubsection{Supervised Methods}
With large-scale paired datasets generated by RGBD cameras, supervised training pipelines of fully-convolutional networks (FCNs) have dominated the SOTA performance over the past decade~\cite{zhao2020monocular}. The prototypical model architectures adopt progressively upsampled feature representations~\cite{fu2018deep,ranftl2020towards}, dilated convolutions~\cite{qiao2021vip}, or parallel multi-scale feature aggregation~\cite{lee2018single,li2017two} to learn fine-grained depth predictions from RGB images. More recent architectures integrate high-resolution representation with multiple lower-resolution \textit{refined} feature maps~\cite{eigen2014depth,eigen2015predicting} to improve the local structural details of the prediction. In recent years, various Vision Transformer (ViT)~\cite{vaswani2017attention}-based spatial \textit{attention blocks}~\cite{bhat2021adabins,ranftl2021vision} are incorporated to ensure global context awareness. Notable objective functions used by these networks are the reverse Huber loss~\cite{zwald2012berhu}, classification and ordinal regression loss~\cite{fu2018deep,cao2017estimating}, pairwise ranking loss~\cite{xian2020structure}, and other adaptive weighting or density losses~\cite{bhat2021adabins,lee2020multi}.

\subsubsection{Semi-supervised and Unsupervised Methods}
The unsupervised learning methods avoid the need for large-scale paired data by inferring scene depth from two or more RGB images by exploiting their inherent geometric constraints~\cite{9803821}. Contemporary methods take stereo image pairs~\cite{garg2016unsupervised, godard2017unsupervised} or consecutive frames from video~\cite{zhou2017unsupervised, godard2019digging} as inputs, combined with geometric view reconstruction~\cite{watson2019self, poggi2018learning} or camera pose estimation~\cite{guizilini20203d} to achieve dense monocular depth prediction. In recent years, attention mechanisms are coupled into such networks to preserve the inherent spatial details of the predicted depth map~\cite{dai2019unsupervised,xu2021multi}. Practical loss functions such as left-right disparity consistency loss~\cite{godard2017unsupervised, yang2018deep}, photometric loss~\cite{zhan2018unsupervised}, and symmetry loss~\cite{zhou2019unsupervised} are generally utilized in these networks as learning objectives. On the other hand, to get more scale information while avoiding costly ground truth, researchers proposed semi-supervised learning methods by combining multiple training stages~\cite{luo2018single, guo2018learning} or training loss terms~\cite{kuznietsov2017semi} of supervised and unsupervised learning. The left-right consistency loss~\cite{amiri2019semi}, mutual distillation loss~\cite{baek2022semi}, and teacher-student learning strategies~\cite{cho2019large} are also introduced into semi-supervised learning pipelines with inspiring results for robust monocular depth estimation.

\subsection{Depth Estimation on Underwater Imagery}
SOTA monocular depth estimation methods are not directly applicable off-the-shelf to underwater imagery due to their domain-specific image formation characteristics and scene geometry~\cite{akkaynak2018revised,drews2016underwater}. Researchers have addressed this in many ways, which can be categorized into active or passive approaches. Prominent \textit{active} approaches use time-correlated single-photon counting on ToF cameras~\cite{maccarone2015underwater} or 3D model adaptation on structure light~\cite{massot2015optical,8559463}. On the other hand, \textit{passive} estimation methods incorporate domain knowledge either following a physics-based IFM or from large-scale data; these models are more popular for their operational simplicity. While atmospheric models were used by early methods~\cite{8084665}, contemporary approaches generally adopt the Jaffe-McGlamery underwater IFM~\cite{6720385} to estimate the depth map from the medium transmission map, \ie, the fraction of scene radiance that reaches the camera after absorption and scattering.

The medium transmission map is generally estimated by using the concepts of dark channel prior (DCP)~\cite{he2010single}, often adapted for underwater scenes as UDCP (underwater DCP)~\cite{7426236}, red-channel compensation~\cite{galdran2015automatic}, etc. Blurriness and illumination information are also used to approximate the transmission map and background light first, and then depth maps are estimated by following the IFM~\cite{8084744,7351749,raihan2020depth,chang2018single}. Hence, depth map estimation from the Jaffe-McGlamery model or following the (more comprehensive) revised IFM~\cite{akkaynak2018revised,akkaynak2019sea} requires accurate estimation of background light, transmission map, and water-body parameters. However, these parameters are not always known; in fact, they are estimated by using noisy depth priors for image enhancement and color correction~\cite{ancuti2017color,skinner2019uwstereonet}.

A practical alternative is to apply transfer learning on SOTA deep visual models for fast monocular depth estimation. The idea is to take advantage of large-scale terrestrial RGB-D datasets for pre-training, and then tune the model weights on limited underwater data. The most commonly adopted models are the Monodepth2~\cite{godard2019digging}, AdaBins~\cite{bhat2021adabins}, various U-Net~\cite{ronneberger2015u} based architectures (\eg, with VGG encoder~\cite{simonyan2014very,laina2016deeper} or ResNet50 encoder~\cite{he2016deep}), and GANs~\cite{hambarde2021uw}. Contemporary researchers have also proposed hybrid training pipelines where in-air (\ie, terrestrial) RGB-D pairs drive the supervised training, while the domain (\ie underwater) data is used in a self-supervised manner simultaneously~\cite{ye2019deep,gupta2019unsupervised}. Despite some early success of these existing approaches, several important aspects of ($i$) incorporating useful domain knowledge into comprehensive training pipelines, ($ii$) learning depth prediction from large-scale underwater RGB-D data, and ($iii$) computational feasibility analysis for robot vision - are not explored in the literature.

\section{Proposed Learning Pipeline}
\subsection{New Input Space Adaptation: RMI}\label{sec_rmi}
Although the wavelength-dependency of atmospheric light attenuation and scattering are negligible, they impact underwater image formation significantly. According to the Jaffe-McGlamery underwater IMF~\cite{6720385}, an observed image $\mathbf U_\lambda(\mathbf x)$ can be expressed as:
\begin{equation}
    \mathbf U_\mathbf \lambda(\mathbf x) = {\mathbf I_\mathbf \lambda(\mathbf x)\cdot t_\mathbf \lambda(\mathbf x)} + {\mathbf B_\mathbf \lambda \cdot (1- t_\mathbf \lambda(\mathbf x))},
    \label{eq4}
\end{equation}
where $\mathbf \lambda \in\{\text{R,G,B}\}$ is the wavelength component, $\mathbf I_\lambda(\mathbf x)$ is the clear latent image, $\mathbf B_\lambda$ denotes the global background light, and $t_\lambda(\mathbf x)$ represents the medium transmission rate. The $t_\mathbf \lambda(\mathbf x)$ is a function of pixel depth $d(x)$ and medium attenuation coefficient $\beta_\mathbf \lambda$, defined as: $t_\mathbf \lambda(\mathbf x) = e^{-\beta_\lambda \cdot d(\mathbf x)}$.

Many contemporary works use priors like DCP~\cite{he2010single} or UDCP~\cite{7426236} to approximate $t_\mathbf \lambda(\mathbf x)$ and estimate coarse scene depth $d(\mathbf x)$. Since red wavelength suffers more aggressive attenuation underwater, techniques such as underwater light attenuation prior (ULAP)~\cite{song2018rapid} and red-channel compensation~\cite{galdran2015automatic} can exploit the R channel values to further refine the depth prediction. As illustrated in Fig.~\ref{fig:rmg_space}, the relative differences between \{R\} and \{G, B\} channel values encode useful depth information for a given pixel. In this paper, we exploit these inherent relationships and demonstrate that R\textbf{M}I$\equiv$\{R, \textbf{M}$=$$max$\{G,B\}, I (intensity)\} is a significantly better input space for visual learning pipelines of underwater monocular depth estimation models.

\vspace{-0.5mm}
\subsection{Network Architecture: UDepth Model}
As illustrated in Fig.~\ref{fig:udepth_arch}, the network architecture of UDepth model consists of three major components: a MobileNetV2-based encoder-decoder backbone, a transformer-based refinement module (mViT), and a convolutional regressor. These components are tied sequentially for the supervised learning of monocular depth estimation.

\subsubsection{\textbf{MobileNetV2 backbone}} We use an encoder-decoder backbone based on MobileNetV2~\cite{sandler2018mobilenetv2} as it is highly efficient and designed for resource-constrained platforms. It is considerably faster than other SOTA alternatives with only a slight compromise in performance, which makes it feasible for robot deployments. It is based on an \textit{inverted residual} structure with residual connections between \textit{bottleneck} layers~\cite{sandler2018mobilenetv2}. The intermediate expansion layers use lightweight depthwise convolutions to filter features as a source of non-linearity. The encoder contains a series of fully convolution layers with $32$ filters, followed by a total of $19$ residual bottleneck layers. We adapt the last convolutional layer of decoder so that it finally generates $48$ filters of $320\times 480$ resolution, given a 3-channel R\textbf{M}I input.

\subsubsection{\textbf{mViT refinement}}
Transformers can perform global statistical analysis on images, solving the problem that traditional convolution models can only handle pixel-level information~\cite{dosovitskiy2020image}. Due to the heavy computational cost of Vision Transformers (ViT), we adopt a lighter mViT architecture inspired by~\cite{bhat2021adabins}. The $48$ filters extracted by the backbone are $1\times1$ convolved and flattened to patch embeddings, which serve as inputs to the mViT encoder. Those are also fed to a $3\times3$ convolutional layer for spatial refinements. The $1 \times 1$ convolutional kernels are subsequently exploited to compute the range-attention maps $\mathbf R$, which combines adaptive global information with local pixel-level information from CNN. The other embedding is propagated to a multilayer perceptron head with ReLU activation to obtain a $80$-dimensional \textit{bin-width} feature vector $\mathbf f_b$.

\begin{figure*}[t]
\centering    \includegraphics[width=0.95\linewidth]{imgs/model.pdf}
\vspace{-3mm}
\caption{The end-to-end learning pipeline of our proposed UDepth model is shown. Raw RGB images are first pre-processed to map into R\textbf{M}I input space, then forwarded to the MobileNetV2 backbone for feature extraction. Those features are refined by a transformer-based optimizer (mViT), followed by a convolutional regressor to generate the single channel depth prediction. The learning objective involving pixel-wise losses and a domain projection loss is formulated in Eq.~\ref{obj_fun}.}
\label{fig:udepth_arch}
\vspace{-5mm}
\end{figure*}

\subsubsection{\textbf{Convolutional regression}}
Finally, the convolutional regression module combines the range-attention maps $\mathbf R$ and features $\mathbf f_b$ to generate the final feature map $\mathbf f$. To avoid discretization of depth values, the final prediction of depth map $\mathbf D$ is computed by the linear combination of bin-width centers ($\bar{\mathbf f_{b}}$), which is given by: $\hat{d} = \sum_{k} \bar{\mathbf f_{bk}} \; \sigma(\mathbf R_k)$.

\subsection{Objective Function Formulation}
\noindent
{\textbf{Pixel-wise depth losses.}} We use two pixel-wise supervised loss functions: $i)$ the $\mathcal{L}_2$ loss, and $ii)$ a scaled version of the Scale-Invariant Log (SILog) loss introduced by Eigen \etal~\cite{eigen2014depth}. These are defined as follows:
\begin{align}
\footnotesize
    \mathcal{L}_{2} &= \mathbb{E}\Big[ \big\lVert {d}_{i} - \hat{d}_{i} \big\rVert_2 \Big],  \\
    \mathcal{L}_{SILog}&=\alpha\sqrt{\frac{1}{T}\sum_{i}{g_i^2-\frac{\lambda}{T^2}(\sum_{i}{g_i})^2}} .
\end{align}
Here, $g_i=\log_{}{\hat{d_i}}-\log_{}{d_i}$ where $\hat{d}$ is the predicted depth, $d$ is the ground truth depth, and $T$ denotes the number of pixels having valid ground truth values. Inspired by~\cite{bhat2021adabins}, we use $\lambda=0.85$ and $\alpha=10$ in our implementation.
\vspace{1mm}

\noindent
{\textbf{Domain projection loss.}\label{sec_dpj}} We formulate a novel domain projection loss function based on our input space calibration. Following our discussion in Sec.~\ref{sec_rmi} and Fig.~\ref{fig:rmg_space}, we express the R-\textbf{M} relationship with the depth of a pixel $x$ with a linear approximator as follows:
\begin{equation}
 \tilde{d}(x \big| R, M) = \mu_0 + \mu_1 R(x) + \mu_2 M(x).
\label{ulap_constraint}
\end{equation}
Then, we find the least-squared solution $\mu^*$ on the entire RGB-D training pairs, which is over $2.8$ billion pixels ($9229$ images of $640\times480$ resolution) by optimizing:
\begin{equation}
 \mu^* = \argmin_\mu \sum_{i,x} {\big\lVert {d}_{i}(x) - \tilde{d}(x \big| R_i, M_i) \big\rVert}_2^2.
\label{rmi_opt}
\end{equation}
We find $\mu^*=[0.496, -0.389, 0.464]$ in our experiments on USOD10K dataset~\cite{USOD10K} (see Sec.~\ref{usod_data}). Here, our goal is to use the optimal $\mu$-space for regularization, we penalize any pixel-wise depth predictions that violate the underwater image attenuation constraint defined by Eq.~\ref{ulap_constraint}. We achieve this by the following projection error function:
\begin{equation}
  \mathcal{L}_{\perp} = \mathbb{E}\Big[\Pi_{\mu^*}\big(\tilde{d}_i(R, M) - \hat{d}_i) \big) \Big]
    \label{perp}
\end{equation}

\noindent
{\textbf{End-to-end objective}.} Finally, the end-to-end learning objective of our proposed UDepth pipeline is formulated as:
\begin{equation}
 \mathcal{L}_{UDepth} = \lambda_2 \; \mathcal{L}_{2} + \lambda_{S}  \;\mathcal{L}_{SILog} + \lambda_{\perp} \; \mathcal{L}_{\perp}.
\label{obj_fun}
\end{equation}
In Eq.~\ref{obj_fun}, we find the $\lambda$-parameters empirically through hyper-parameters tuning. The optimal values used in UDepth training are: $\lambda_2 = 0.3$, $\lambda_{S} = 0.6$, and $\lambda_{\perp} = 0.1$.
\section{Conclusion}
Fast monocular depth estimation can facilitate real-time 3D perception capabilities of autonomous underwater robots. In this work, we propose a deep supervised learning pipeline that includes: ($i$) a domain-aware input space adaptation based on underwater light attenuation characteristics of light propagation; ($ii$) a least-squared formulation of the attenuation constraints for domain projection of coarse underwater scene depth; and ($iii$) an efficient deep visual model named UDepth, which can be trained by that domain projection loss and other pixel-level losses for fine-grained monocular depth estimation. The UDepth model is designed with MobileNetV2 backbone and a Transformer-based optimizer to be able to learn an efficient and robust solution for embedded devices. Experimental results show that with only $20\%$ computational cost, UDepth offers comparable and often better depth estimation performance than SOTA models on benchmark datasets and arbitrary test cases. The domain projection with additional filtering also provides a fast solution for low-powered robots. In the future, we plan to extend UDepth's capabilities toward self-supervised depth estimation and depth-guided image enhancement for real-time underwater robot vision.


\section{Introduction}

The problem of the Galactic halo is being discussed from the beginning of the 1950s. Before that time, a sharp transition
was assumed between the Galactic disk of the thickness of $\sim100$~pc and the extragalactic medium. Later, \citet{ginz53}
developed conceptions of the physical cosmic ray (CR) halo with the size of about 15~kpc, where CRs are trapped by
scattering (i.e., propagate diffusively). The characteristic CR age in the Galaxy was estimated as $\sim 10^8$~yr, which is
confirmed by radio data and by the information on CR chemical composition (e.g., about the abundance of unstable isotope
$^{10}$Be) \citep[see, e.g.,][etc.]{Ginzburg1976,sza80}.

The static halo model of a fixed height was presented in \citet{syr64}. It assumes that the CR density at a certain distance
from the Galactic plane becomes negligible. This model is currently broadly implemented in advanced numerical codes, such as
GALPROP \citep{mosk98}.

The downside of the model is that it depends on two arbitrary parameters, namely the diffusion coefficient and the halo
size, whose values are ambiguously defined. Therefore, it is necessary to describe the processes of generation and damping
of MHD turbulence in the halo and their connections to the CR transport self-consistently.

In \citet{dogiel2020}, we have suggested a model of CR self-confinement in the Galaxy, where the turbulence generated in the
Galactic disk was amplified by streaming CRs. However, the turbulence excitation rate is very high in that model, and hence
the size of the halo is too small at GeV energies. To resolve this issue, in the present paper we also take into account the
nonlinear Landau (NL) damping, which was neglected in the original work. We show that the inclusion of the damping term
leads to a significantly larger halo size. We also show that the MHD turbulence which confines CRs in the halo can be
entirely self-generated by CRs.

\section{Self-consistent nonlinear Model of CR halo}

Unlike models with a pre-defined halo size, self-consistent halo models include a mechanism of MHD-wave excitation. In these
models, CR propagation is described by a system of nonlinear equations \citep[see, e.g., in][and references
therein]{dogiel:94,evoli18,dogiel2020}.

A general system of simplified one-dimensional nonlinear equations for the CR spectrum $N(p,z)$ and the energy density of
MHD fluctuations $W(k,z)$ can be presented in the following form:
\begin{equation}
\begin{array}{l}
{\displaystyle\frac{\partial}{\partial z} \left(u_{\rm adv}N  - D\frac{\partial N}{\partial z}\right)
-\frac{\partial}{\partial p}
\left(\frac13\frac{du_{\rm adv}}{dz}pN -\dot{p}N\right)}
= Q \,,\vspace{.2cm}\\
{\displaystyle\frac{\partial u_{\rm A}W}{\partial z} - \frac{du_{\rm A}}{dz}\frac{\partial (kW)}{\partial k}
+\frac{\partial}{\partial k}\left(\frac{kW}{\tau_{\rm cas}}\right)}=
(\Gamma_{\rm CR} - \nu)W \,,
\end{array}
\label{system}
\end{equation}
where $Q(p,z)$ is the source term of CRs, $u_{\rm adv}(z)$ is the CR advection velocity, which depends on the difference
between outward- and inward-propagating MHD waves, $u_{\rm A}(z)=B(z)/\sqrt{4\pi\rho(z)}$ is the Alfven velocity, $\dot{p} <
0$ is the rate of momentum loss due to interaction with gas, $\rho = m_p n$ is the mass density of ionized hydrogen ($m_p$
is proton mass), and $B$ is the strength of the longitudinal large-scale magnetic field. Furthermore, $\Gamma_{\rm CR}(k,z)$
is the rate of resonant wave excitation, $\nu$ is the wave damping rate, and $\tau_{\rm cas}(W)$ is the characteristic
timescale of turbulent cascade to larger $k$; the latter depends on the particular process of MHD-generation \citep[see,
e.g.,][this process is discussed in Section~\ref{sec:mirror}]{ptus06}. The spectrum $N(p,z)$ is normalized such that $\int
N(p) dp$ is the total number density of CRs.

The wavenumber $k$ of MHD fluctuations is related to the CR momentum $p$ via the resonance condition \citep{skil3},
\begin{equation}\label{resonant}
kp\approx m_p\Omega_*\,,
\end{equation}
where $\Omega_*=eB/m_pc$ is the gyrofrequency of non-relativistic CR protons. The resulting CR diffusion coefficient is
\citep{skil3}
\begin{equation}
D(p,z)\approx\frac{vB^2}{6\pi^2k^2W}\,.
\label{ddif}
\end{equation}
In this approximation, the excitation rate is proportional to the CR diffusion flux,
\begin{equation}
\Gamma_{\rm CR}(k,z)\approx-\frac{2\pi^2eu_{\rm A}p}{Bc}\,D\frac{\partial N}{\partial z}\,.
\label{gmm}
\end{equation}

There are very few known parameters and processes that can govern the density of MHD-fluctuations in the halo (and thus the
CR diffusion). These are the spatial dependencies of the magnetic field and the gas density, the magnitude and the spectrum
of CR source in the disk, and nonlinear processes involving MHD waves. In this respect, the variety of models for the wave
excitation in the halo is very restricted.

\subsection{Development of \citet{dogiel2020}} \label{sec:mirror}

\citet{evoli18} and \citet{dogiel2020} presented one-dimensional models of CR propagation along the magnetic field lines,
with MHD-fluctuations excited by the resonant CR-streaming instability.

\citet{evoli18} developed a model of MHD-turbulence with nonlinear cascading to larger $k$. They considered three sources
of waves responsible for CR scattering in the halo: (i) self-generated MHD-waves excited by CRs through the streaming
instability, (ii) processes mimicking wave generation by, e.g., supernova explosions in the disk which eject waves at large
scales, and (iii) cascading process which is determined by an initial arbitrary source of background turbulence distributed
over the halo. In case the cascading is responsible for the damping of MHD-fluctuations in the halo, the CR halo can be
about of few kpc, which is compatible with the estimations of GALPROP.

On the contrary, \citet{dogiel2020} showed that the cascading process in the halo is negligible for relevant values of $k$,
i.e., the term containing $\tau_{\rm cas}$ in the second Equation~(\ref{system}) can be omitted. We considered two sources
of waves responsible for CR scattering in the halo, namely, (i) self-generated MHD waves excited by CRs through the
streaming instability, and (ii) the spectrum of MHD-fluctuations generated by sources in the Galactic disk. In that model,
magnetic fluctuations are only excited in the direction away from the disk. However, the resulting CR flux excites waves too
efficiently, which yields the halo size of only $\sim100$~pc at low energies.

In \citet{dogiel2020}, we have not considered the possibility that outgoing waves may be reflected by a nonuniform medium
\citep[see][for details]{ferraro54,kulsrud2005}. In fact, that happens if the approximation of geometrical optics is no
longer applicable. According to \citet{Ginzburg1970}, if the wave phase velocity changes from $u_{\rm min}$ to $u_{\rm max}$
within a layer of thickness $\ell$, the reflection coefficient $R$ of the outgoing waves from the layer is
\begin{equation}
R^2 \sim \exp \left( - 4\pi k\ell\frac{u_{\rm min}}{u_{\rm max}} \right) \,.
\end{equation}
Applying this expression to the halo with $\ell = 1$ kpc, we see that even for very long waves, resonant with PeV protons,
only $\sim0.1\%$ of the total energy is reflected. Therefore, we indeed can safely assume that there are no backward-moving
waves present in the halo.

Another physical process neglected in \citet{dogiel2020} is nonlinear Landau (NL) damping. A two-dimensional halo
model including this process has already been developed by \citet{dogiel:93} and \citet{dogiel:94}. The authors
used the equation for CR propagation complemented with the equation for MHD-fluctuations which are excited by the CR flux
and attenuated by the NL damping, cascading, and adiabatic losses. It was shown that CR distribution is quasi-isotropic near
the Galactic plane, but becomes more focused along the radial coordinate as particles propagate further away. At some point
the scattering becomes unable to reflect particles back, which sets the outer boundary of the halo. The halo size was
estimated to be about 10~kpc. However the advective transport of CRs was not taken into account in this model.

According to \citet{voelk1982} and \citet{miller1991}, the rate of NL damping is given by
\begin{equation}
\nu_{\rm NL}(k) \approx g(n,T)\frac{8\pi u_{\rm A}}{B^2} k\int\limits_{k_{\rm min}}^k W(k_1) dk_1 \,,
\label{NLD}
\end{equation}
where the lower integration limit $k_{\rm min}$ is unimportant for our self-consistent model (see
Section~\ref{sec:Analytic_mid}). Following \citet{miller1991}, the dimensionless factor $g(\beta)$ in Equation~(\ref{NLD})
can be approximated by
\begin{equation}
g(\beta) \approx \frac{\sqrt{\pi}}{4}\beta^{1/2}\left(e^{-\beta^{-1}} + \frac{1}{2}\epsilon^{1/2}e^{-\epsilon\beta^{-1}}\right) \,.
\label{eq:g_approximation}
\end{equation}
Here, $\epsilon = m_e/m_p$ is electron-to-proton mass ratio and
\begin{equation}
\beta=\frac{nkT}{B^2/8\pi} \equiv \frac{u^2_{\rm th}}{u_{\rm A}^2}\,,
\end{equation}
is the plasma-$\beta$ parameter expressed via the thermal velocity of protons $u_{\rm th}$.
We assume that temperatures $T$ of both protons and electrons are equal.

The idea that the excitation of MHD turbulence in a halo can be balanced by NL damping has been previously discussed by
\citet{ptuskin97}. In the present work, we use a different expression for NL damping, which takes into account contributions
of both thermal protons and electrons. Furthermore, in our model $\beta$ significantly drops with the height, which results
in a much weaker damping for waves excited by CRs with energies above 100 GeV and, thus, helps to confine such particles.

\section{The halo model with NL damping}

\begin{figure}
\begin{center}
	\includegraphics[width=\columnwidth]{schema.eps}
    \caption{Schematic representation of the model considered in the paper.}
    \label{fig:schema}
\end{center}
\end{figure}

The idealized structure of our model is sketched in Figure~\ref{fig:schema}. We consider two characteristic regions along
the $z$-axis: the Galactic disk, where MHD-turbulence is assumed to be generated by sources distributed over the Galactic
plane, and the CR halo, where the turbulence is self-generated by the outgoing CR flux. We assume that the magnetic field is
practically vertical in the halo (while its geometry can be arbitrary in the disk), and that CRs do not diffuse across the
magnetic field lines.

This model allows us to reduce the set of equations~(\ref{system}) to
\begin{eqnarray}
\frac{\partial}{\partial z} \left(
u_{\rm adv}N - D\frac{\partial N}{\partial z}\right) -
\frac{\partial}{\partial p}\left(\frac{1}{3}\frac{du_{\rm adv}}{dz} pN -\dot{p}N\right)\label{eq:general_protons}\hspace{.7cm}\\
= 2Q(p)\delta(z) \,,\nonumber\\
\frac{\partial v_{\rm A}W}{\partial z} - \frac{du_{\rm A}}{dz}\frac{\partial kW}{\partial k} =
(\Gamma_{\rm CR} - \nu_{\rm NL})W \,,\hspace{2.2cm}
\label{eq:general_waves}
\end{eqnarray}
where $Q(p)$ is the source of CRs above/below the Galactic plane, $\dot{p} < 0$ is the momentum loss rate due to ionization
or proton-proton collisions within the disk ($0 < z < z_d$), while in the halo ($z\geq z_d$) the CR losses are solely due to
adiabatic cooling. Here and below, $z_d$ is the characteristic height of the disk. The CR advection velocity changes
discontinuously at the disk boundary,
\begin{equation}
u_{\rm adv}(z) = u_{\rm A}\theta(z - z_d) \,,
\label{eq:disk_to_halo_conv}
\end{equation}
where $\theta(z)$ is the Heaviside function. We assume that sources of the turbulence in the disk do not contribute to the
turbulence in the halo, i.e.
\begin{equation}
W(k, z_d) = 0 \,,
\end{equation}
i.e., the halo turbulence is entirely self-generated by CRs. However, the existence of turbulence within the disk is
essential \citep[][]{evoli18,dogiel2020}, and we take this into account in Section~\ref{sec:Analytic_strt}.

According to Equation~(\ref{eq:g_approximation}), MHD-waves are damped on plasma electrons in a low-$\beta$ plasma, and on
protons in a high-$\beta$ plasma. Most of the thermal electrons or protons contribute to the damping if, respectively,
\begin{equation}
0.01 \leq \beta \leq 0.2,
\label{range1}
\end{equation}
or
\begin{equation}
\beta \geq 10 \,.
\label{range2}
\end{equation}
For these values of $\beta$, the respective dominant exponential factor in Equation~(\ref{eq:g_approximation}) can be set to
unity, and $\nu_{\rm NL}$ in Equation~(\ref{NLD}) becomes independent of the plasma density.

Below in Section~\ref{sec:Analytic_strt} we obtain a simplified analytical solution of Equations~(\ref{eq:general_protons})
and (\ref{eq:general_waves}), while in Section~\ref{sec:num_sol} we present and discuss the exact numerical solution.

\section{Analytic approximation} \label{sec:Analytic_strt}

In this section we derive an analytical solution which qualitatively explains the role of NL damping in the self-consistent
halo model.

\subsection{CR spectrum in the disk, and CR flux from the disk to the halo}

To simplify CR propagation in the disk \citep[see, e.g.][]{ber90}, we consider the following two key parameters: the
outgoing CR flux and CR distribution function at the boundary between the disk and the halo. They both should be continuous
at boundary.

A general equation for the outgoing flux $S_0(p) = S(p,z_d)$ at the boundary is obtained by integrating
Equation~(\ref{eq:general_protons}),
\begin{equation}
S_0(p) = Q(p) + \frac{d}{dp}\left(\frac{1}{3}u_{\rm A0}\, pN_0(p)-\int\limits_0^{z_d}\dot{p}N(z,p)dz\right) \,,
\label{eq:jump_condition}
\end{equation}
where $u_{\rm A0} = u_{\rm A}(z_d)$ and $N_0(p) = N(p,z_d)$. The flux $S_0$ derived from Equation~(\ref{eq:jump_condition})
can be considered as the boundary condition for Equation~(\ref{eq:general_protons}) at $z=z_d$. Since the halo size should
be much larger than $z_d$, we assume that the CR spectrum does not change significantly across the disk. Then we can
approximate $N(p,z) \approx N_0(p)$ for $0 \leq z \leq z_d$.

As pointed out in \citet{dogiel2020}, energy losses in the halo are unimportant and thus the CR flux $S_0(p)$ is conserved.
Then we obtain the following solution of Equation~(\ref{eq:general_protons}) for $z \geq z_d$:
\begin{equation}
N(p,z) = \frac{S_0(p)}{u(p,z)}\,,
\label{eq:f_general_solution}
\end{equation}
where $u(p,z)$ is the outflow velocity of CRs,
\begin{equation}
u = \left(\int\limits_\eta^{\eta_\infty} \frac{e^{\eta-\eta_1}d\eta_1}{u_{\rm A}(\eta_1)} \right)^{-1} \,,
\end{equation}
and $\eta(p,z)$ is a dimensionless variable
\begin{equation}
\eta = \int\limits_{z_d}^z \frac{u_{\rm A}}{D} dz_1 \,.
\label{eq:eta}
\end{equation}
The value of $\eta_\infty$ generally depends on the boundary condition at $z \rightarrow \infty$. Substituting $N_0(p)=N(p,
z_d)$ from Equation~(\ref{eq:f_general_solution}) in Equation~(\ref{eq:jump_condition}), we derive the flux,
\begin{equation}
S_0(p) = \frac{u_d(p)}{\mathcal{E}(p)}\int\limits_p^\infty Q(p_1) \exp \left( - \int\limits_p^{p_1}
\frac{u_d(p_2)}{\mathcal{E}(p_2)} dp_2 \right) dp_1 \,.
\label{eq:jump_expression}
\end{equation}
Here, $u_d(p) = u(p,z_d)$ and $\mathcal{E}(p) = \frac13pu_{\rm A0}-\int_0^{z_d}\dot{p}dz = \frac13pu_{\rm A0} +
\frac12\mathcal{N}_{\rm H} L(p)$, where $\mathcal{N}_{\rm H}$ is the vertical column density of hydrogen atoms in the disk
and $L(p) = -\dot{p}/n_{\rm H}$ is energy loss function (per unit column density) due to interaction with the disk gas. We
note that $u_d$ in fact depends on $S_0$, and, therefore, Equation~(\ref{eq:jump_expression}) is an integral equation for
$S_0(p)$. If the dependence $u_d$ versus $S_0$ is weak, the equation can be solved
iteratively.

We can obtain a simple approximation for $S_0(p)$ assuming that $\mathcal{E}S_0/u_d$ is a power-law function $\propto
p^{-\alpha}$. According to experimental data, $S_0/u_d\equiv N(p)$ has a negative spectral index smaller than that of
$S_0(p)$ (both are smaller than $-2$), while $\mathcal{E}$ cannot increase faster than $\propto p$. Therefore, $\alpha>0$
and we readily obtain from Equation~(\ref{eq:jump_condition}),
\begin{equation}
S_0(p) = \frac{Q(p)}{1 + \alpha\mathcal{E}/(pu_d)} \,.
\end{equation}
We notice that both the energy loss rate $\mathcal{E}/p$ and the inverse outflow velocity $u_d^{-1}=N/S_0$ decrease with
$p$, and therefore $S_0(p) \approx Q(p)$ for sufficiently high energies. To evaluate a critical energy where
$\mathcal{E}/(pu_d)=1$, we assume $\mathcal{N}_{\rm H} \approx 6\times 10^{20}$~cm$^{-2}$ and $u_d \sim u_{\rm A}\sim
10^6$~cm/s. This yields the proton energy about 0.5~GeV, above which we can set $S_0 = Q$.

As discussed in \citet{dogiel2020}, the value of $\eta$ is the key parameter characterizing CR propagation in the halo. The
entire halo can be approximately split into two regions: one is called the {\it halo sheath}, where $\eta(z) \ll 1$ and the
diffusion term $-D\,\partial N/\partial z$ dominates in the CR flux $S_0$; the other is where $\eta(z) \gg 1$ and the
advection term $u_{\rm adv}N$ dominates. The critical point $z_{\rm cr}(p)$ separating these two regions can be determined
from the condition $\eta(p,z_{\rm cr})\approx1$. Since the dominant advection irreversibly carries CRs away from the disk,
the halo size can be set equal to $z_{\rm cr}$. Therefore, the boundary condition at $z \rightarrow \infty$ becomes
unimportant as long as $\eta_\infty \gg 1$.

In order to derive $\eta$ from Equation~(\ref{eq:eta}), we need to obtain the diffusion coefficient $D$ from
Equation~(\ref{ddif}), which requires the solution of Equation~(\ref{eq:general_waves}).

\subsection{Excitation-damping balance}\label{sec:Analytic_mid}

The numerical solution of Equations~(\ref{eq:general_protons}) and (\ref{eq:general_waves}) (see
Section~\ref{sec:num_sol}) suggests that $W(k,z)$ in the diffusion region can be estimated from the excitation-damping
balance,
\begin{equation}
\Gamma_{\rm CR} = \nu_{\rm NL}\,.
\end{equation}
We rewrite it using Equations~(\ref{gmm}) and (\ref{NLD}),
\begin{equation}
\frac{4g(z)c^2}{\pi e^2B^2} k^2\int \limits_{k_{\rm min}}^k W(k_1) dk_1 = S_0(p) - u_{\rm A}N\,.
\label{gz}
\end{equation}
In the halo sheath ($\eta < 1$) the last term of Equation~(\ref{gz}) can be neglected. In this case, $S_0(p)\propto Q(p)$
decreases with $p$ faster than $p^{-2}$, and thus the integral on the left-hand size of Equation~(\ref{gz}) is dominated by
the upper limit $k$. Therefore,
\begin{equation}
W(k,z) = \frac{\pi}{4g(z)}\frac{\partial}{\partial k}\left[p^2S_0(p)\right] \,,
\end{equation}
and
\begin{equation}
\eta(p,z) = -\frac{3\pi^3e}{2vc}\frac{\partial}{\partial p}\left[p^2S_0(p)\right] \int\limits_{z_d}^{z}
\frac{u_{\rm A}}{Bg(z_1)}dz_1 \,.
\end{equation}
From  Equation~(\ref{eq:g_approximation}) we obtain
\begin{equation}
\eta = -\frac{6\pi^{5/2}e}{vc}\frac{\partial}{\partial p}\left[p^2S_0(p)\right] \int\limits_{z_d}^{z}
\frac{u_{\rm th}(T)\beta^{-1}dz_1}{B(e^{-\beta^{-1}} + \frac12\epsilon^{1/2}e^{-\epsilon\beta^{-1}})} \,.
\label{eq:eta_notexpanded}
\end{equation}
For simplicity, below we assume $n(z)=n_0 \exp(-z/z_n)$, $B(z)=B_0 \exp(-z/z_B)$, and $T(z)=T_0 \exp(z/z_T)$.

\subsection{Spectrum of CRs in the halo sheath}
\label{sec:Analytic_end}

If $\beta$ is within the ranges defined in Equations~(\ref{range1}) and (\ref{range2}), the expression for $g(\beta)$
simplifies significantly. In this regime, previously considered by \citet{ptuskin97},  we can neglect the exponential
dependence in the denominator of Equation~(\ref{eq:eta_notexpanded}) and rewrite the equation as
\begin{equation}
\eta(p,z) = A(p) \left(e^{z/z_\eta} - e^{z_d/z_\eta}\right)\,,
\label{eq:eta_const_damping}
\end{equation}
where $z_\eta^{-1} = z_n^{-1} - z_B^{-1} - \frac12z_T^{-1}$. For the sake of simplicity, below we assume $z_d \approx 0$.

The magnitude of the dimensionless factor $A(p)$ depends on the dominant mechanism of NL damping. In a low-$\beta$ plasma
with $0.01 \leq \beta \leq 0.2$ the damping on thermal electrons dominates, and
\begin{equation}
A(p) \approx A_e(p) = -\frac{12\pi^{5/2}eu_{\rm A0}^2z_\eta}{vcB_0u_{\rm th0}\epsilon^{1/2}}\,\frac{\partial}{\partial p}
\left[p^2S_0(p)\right] \,,
\label{eq:alow_beta}
\end{equation}
where $u_{\rm th0} = u_{\rm th}(z_d)$, while in a high-$\beta$ plasma with $\beta \geq 10$ the damping is due to thermal
protons, and
\begin{equation}
A(p) \approx A_p(p) =  \frac12\epsilon^{1/2}A_e(p)\,.
\label{ahp}
\end{equation}
The critical point $z_{\rm cr}$ is derived from the condition $\eta \approx 1$. Thus, the halo size is estimated from
Equation~(\ref{eq:eta_const_damping}) as
\begin{equation}
z_{\rm cr}(p) = z_\eta \ln \left[1 + 1/A(p) \right] \,.
\label{z_cr}
\end{equation}
For low energies, where $A(p) \gg 1$, the halo size increases with $p$ as $z_{\rm cr}(p)\approx z_\eta/A(p)$; for high
energies, the halo size $z_{\rm cr}(p) \approx -z_\eta \ln A(p)$ is almost independent of $p$. We point out that the
model is not viable in the former regime, normally corresponding to the electron-dominated damping, because the resulting
halo size becomes too small. On the other hand, for the proton-dominated damping with $\beta > 10$, the function $A_p(p)\sim
1$ for the following halo parameters:  $B \leq 1~\mu$G, $n \geq 10^{-2}$~cm$^{-3}$, and $T \geq 100$~eV. The halo size in
this case exceeds 1~kpc at energies above 1~GeV.

The CR spectrum is given by Equation~(\ref{eq:f_general_solution}),
\begin{equation}
N(p,z) = \frac{S_0(p)}{u_{\rm A0}}\int\limits_\eta^\infty \frac{e^{\eta-\eta_1}d\eta_1}{[1 + \eta_1/A(p)]^{z_\eta / z_{\rm A}}} \,,
\end{equation}
where $z_{\rm A}^{-1} = \frac12z_n^{-1} - z_B^{-1}$ characterizes the spatial scale of variation of $v_{\rm A}(z)$. This
result can be expressed in terms of the incomplete gamma-function $\mathbf{\Gamma}(a,z)$,
\begin{eqnarray}
N(p,z) = \frac{S_0(p)}{u_{\rm A0}}e^{\eta + A(p)}A(p)^{z_\eta/z_{\rm A}}\hspace{2cm} \\
\times\mathbf{\Gamma}(1 - z_\eta/z_{\rm A}, A(p) + \eta) \,.\nonumber
\end{eqnarray}
If $A(p) + \eta \gg 1$, the solution corresponds to the advection flux with the Alfven velocity at the halo periphery. This
represents low-energy CRs at large distances from the disk,
\begin{eqnarray}
N(p,z) \approx \frac{S_0(p)}{u_{\rm A0}}[1 + \eta/A(p)]^{-z_\eta / z_{\rm A}}\hspace{2cm} \label{eq:f_const_low}\\
= \frac{S_0(p)}{u_{\rm A0}}e^{-z / z_{\rm A}}\equiv \frac{S_0(p)}{u_{\rm A}(z)} \,.\nonumber
\end{eqnarray}
If $A(p) + \eta \ll 1$, the CR spectrum tends to
\begin{eqnarray}
N(p,z) \approx \frac{S_0(p)}{u_{\rm A0}} A(p)^{z_\eta/z_{\rm A}} \mathbf{\Gamma}(1 - z_\eta/z_{\rm A})\hspace{1.6cm}
\label{eq:f_diffusion_appx}\\
\propto S_0(p)[pS_0(p)]^{z_\eta/z_{\rm A}}\,.\nonumber
\end{eqnarray}
The resulting spectrum, corresponding to the diffusion-dominated flux, does not practically depend on $z$ up to the critical
point $z_{\rm cr}$.  We note that Equation (\ref{eq:f_diffusion_appx}) resembles Equation~(34) from \citet{ptuskin97}.

The derived approximate solution has important implications. We conclude that CRs escape from the halo at $z=1-10$~kpc, and
for energetic CRs the halo size weakly depends on their energy. Given $S_0(p) \propto p^{-2.4}$ and $N(p) \propto p^{-2.7}$,
our solution suggests that $z_\eta/z_{\rm A} \approx 0.3/1.4$ or $z_n/z_B - 0.15z_n/z_T = 0.35$.

\section{Numerical solution and discussion}\label{sec:num_sol}

Equations~(\ref{eq:f_const_low}) and (\ref{eq:f_diffusion_appx}) provide sufficiently good approximations for the CR
spectrum as long as $\beta$ is within the ranges defined in Equations~(\ref{range1}) and (\ref{range2}). However, the
magnitude of $\beta$ varies strongly with $z$ and, therefore, the exponential terms in the denominator of
Equation~(\ref{eq:eta_notexpanded}) cannot be generally ignored. As a result, the expressions for $\eta$ and $N$ become
complicated and can only be obtained numerically.

To reduce the number of free parameters, we consider a simple isothermal model ($z_T^{-1}=0$) with a constant magnetic field
($z_B^{-1} =0$). We use the following set of parameters: $B = 1~\mu$G, $n_0 = 0.1$~cm$^{-3}$, $z_n = 1$~kpc, and $T =
10$~eV. For the CR source function, we use $Q(p) \simeq Q_* (p/m_pc)^{-2.4}$ with $Q_*m_pc = 9.4\times
10^{-4}$~cm$^{-2}$~s$^{-1}$, which is similar to the value given by \citet{strong10}. In this case, $\beta = 40$ at $z = 0$.
The total power of CR sources in the Galaxy can be roughly estimated as $\mathcal{W} = 2\pi R_{\rm Gal}^2\int E_{\rm
kin}(p)Q(p) dp = 4.5\times 10^{-3}~\mbox{erg}~\times 2\pi R_{\rm Gal}^2Q_*m_pc$. Assuming $R_{\rm Gal} = 20$ kpc for the
Galactic disk radius, we obtain $\mathcal{W} \approx 10^{41}$ erg/s.

Equations~(\ref{eq:general_protons}) and (\ref{eq:general_waves}) are solved numerically by employing the procedure
described in \citet{dogiel2020}. To account for CR species heavier than protons, the excitation rate $\Gamma_{\rm CR}$ is
multiplied by a factor of 1.5 \citep[see, e.g.,][]{dogiel2018}. For the initial CR density we use $N(p,z) = 0$, to avoid
appearance of a sharp discontinuity at the upper halo boundary. Both the initial MHD spectrum and the boundary condition at $z = 0$
are equal to a small non-zero function $W_0(k,z)$, as it is necessary for the waves excitation. To ensure a weak (logarithmic)
dependence of the integral in Equation~(\ref{NLD}) on the integration limits, we use $W_0(k,z) \propto k^{-1}$.

The resulting halo size, $z_{\rm cr}$, and the differential CR spectrum, $N/4\pi$, are plotted versus the proton kinetic
energy $E_{\rm kin}$ by the dashed lines in Figures~\ref{fig:size_halo_comparison} and \ref{fig:CR_spectra},
respectively. To account for the solar modulation, we use the force-field approximation with potential $\phi = 0.5$ GV
\citep[][]{gleeson68}. Observational data are taken from \citet{agu15} (AMS-02), \citet{calet19} (CALET), \citet{nucleon19}
(NUCLEON), \citet{cream11} (CREAM-I), \citet{cream17} (CREAM-I+III), and \citet{an2019} (DAMPE). The data are collected
using Cosmic-Ray DataBase (CRDB v4.0) by \citet{Maurin2020}.

\begin{figure}
\begin{center}
	\includegraphics[width=\columnwidth]{halo_sizes_92_94_num.eps}
    \caption{Halo size $z_{\rm cr}(E_{\rm kin})$ obtained from the numerical solution of our model. The dashed line shows the case of
    a single-component gas with $B = 1~\mu$G, $n_0 = 0.1$~cm$^{-3}$, $z_n = 1$~kpc, and $T = 10$~eV ($\beta = 40$),
    the solid line represents the case of a two-component gas (see Section~\ref{sec:num_sol}).}
    \label{fig:size_halo_comparison}
\end{center}
\end{figure}

\begin{figure}
\begin{center}
	\includegraphics[width=\columnwidth]{CR_NLD_solution92_94_num.eps}
    \caption{Energy spectra of CR protons obtained from the numerical solution of our model (lines) and the observational data
    (symbols). All parameters are the same as in Figure~\ref{fig:size_halo_comparison}.}
    \label{fig:CR_spectra}
\end{center}
\end{figure}

We stress that the CR spectra strongly depend on a particular model of NL damping. In our case, the damping is
described by Equations~(\ref{NLD}) and (\ref{eq:g_approximation}). Since $\beta$ rapidly drops with the height, so does the
damping and, hence, the CR diffusion coefficient. Therefore the CR spectra plotted in Figure~\ref{fig:CR_spectra} can be
interpreted as follows:
\begin{itemize}
\item $E_{\rm kin} < 10$ GeV: At such energies, $A(p)$ is sufficiently large and, thus, the halo size is small in
    accordance with Equation~(\ref{z_cr}). For this reason, $\beta(z_{\rm cr})\approx \beta(0)>10$ and NL damping is due
    to thermal protons. Equation~(\ref{eq:f_const_low}) is applicable, which gives $N(p) \propto Q(p) \propto p^{-2.4}$.
\item $10~\mbox{GeV} < E_{\rm kin} < 1~\mbox{TeV}$: $N(p)$ starts approaching a softer spectrum described by
    Equation~(\ref{eq:f_diffusion_appx}). The halo size increases with energy as $z_{\rm cr}(p)\propto 1/A(p)$, and thus
    $\beta(z_{\rm cr})$ rapidly decreases, so that eventually a mixed damping both on thermal protons and electrons
    operates.
\item $100~\mbox{GeV} < E_{\rm kin} < 10~\mbox{TeV}$: In the mixed-damping regime, a smooth transition from $A_p(p)$ to
    much larger $A_e(p)$ occurs. According to Equation~(\ref{eq:f_diffusion_appx}), that makes $N(p)$ harder
		(NL damping rapidly reduces with CR energy as the proton contribution becomes negligible,
		and therefore the CR confinement increases). In Figure~\ref{fig:CR_spectra}, the transition is manifested by
the increase seen at $1~\mbox{TeV} < E_{\rm kin} < 10~\mbox{TeV}$.
\item $E_{\rm kin} > 10~\mbox{TeV}$: Finally, at very high energies $\beta(z_{\rm cr})$ decreases below $0.1$, where the
    damping is due to thermal electrons. Equation~(\ref{eq:f_diffusion_appx}) becomes applicable; since $z_\eta / z_{\rm
    A} = 1/2$ in our case, $N(p) \propto p^{-3.1}$.
\end{itemize}

Figure~\ref{fig:CR_spectra} shows that the theoretical curve and the experimental data are in good qualitative agreement.
However, we should also keep in mind that gas in the halo consists of several components. In particular, the warm ionized
gas (WIM) dominates at lower altitudes, while at higher $z$ it is mostly hot coronal gas \citep{ferriere1998,gaen08}. To
account for multiple gas components, we assume that the total gas density in our model is determined by a sum of the two
phases: $n(z) = n_{\rm hot}(z) + n_{\rm WIM}(z)$. The same principle applies to the magnitude of NL damping in
Equation~(\ref{NLD}): $g(z) = g(\beta_{\rm hot}) + g(\beta_{\rm WIM})$. Note that the factor $ku_{\rm A}$ in
Equation~(\ref{NLD}) is the wave frequency, and therefore is the same in both phases. Assuming $B = 1~\mu$G, we use the
following set of parameters:
\begin{itemize}
\item Warm phase ($\beta = 4$): $n_0 = 0.1$~cm$^{-3}$, $T = 1$~eV, $z_n = 0.4$~kpc.
\item Hot phase ($\beta = 4$): $n_0 = 10^{-3}$~cm$^{-3}$, $T = 100$~eV, $z_n = 2$~kpc.
\end{itemize}
The source function is $Q(p) \simeq Q_*(p/m_pc)^{-2.32}$ with $Q_*m_pc = 9\times 10^{-4}$~cm$^{-2}$~s$^{-1}$.

The results for the two-phase model are depicted in Figures~\ref{fig:size_halo_comparison} and \ref{fig:CR_spectra} by the
solid lines. We see that the theoretical curve show a much better agreement with the observational data in this case, which
is due to a much weaker dependence of $\beta$ on $z$.

While the two-phase model provides a remarkably good overall agreement with the experimental data in a wide energy range,
the discrepancy below 10 GeV is up to $20\%$. We believe that this is because the effect of disk turbulence on the vertical
profile of the spectrum can no longer be ignored at such low energies. Indeed, by deriving
Equation~(\ref{eq:jump_expression}) we assume that $N(p,z_d) = N(p,0)$. While this assumption is certainly reasonable for
high energies, the low-energy part of the spectrum should be stronger affected by the fact that the diffusion coefficient in
the disk decreases with energy, which inevitably leads to an increasing vertical gradient of $N(z)$. Therefore, the
low-energy spectrum should be more inhomogeneous at $0 < z < z_d$, and the actual spectrum at $z = 0$ should go somewhat
above the theoretical curves plotted in Figure~\ref{fig:CR_spectra}. Furthermore, the diffusion coefficient in the
Galactic disk is likely not affected by the CR streaming (e.g., due to heavy damping on neutrals), but rather depends on
external sources of turbulence (such as supernova explosions and stellar winds).

Apart from the halo size, another important parameter characterizing propagation of CRs is their grammage $X$, i.e.,
the average surface density traversed by CRs during their lifetime in the Galaxy. The grammage determines the ratio of
secondary-to-primary nuclei, and thus can be derived from experimental data. For our model, it can be roughly estimated as
\begin{equation}
X \approx \mathcal{N}_{\rm H}m_p \frac{c}{u_d} \,,
\end{equation}
which gives $X(10~\mbox{GeV}) \approx 12$ g/cm$^2$ for our parameters at $E_{\rm kin} = 10$~GeV. This value is close to that
obtained by, e.g., \citet{engelmann90}. We stress, however, that such estimates are very approximate: to properly test the
model, we need to accurately calculate the spectra of secondary and primary nuclei, and compare them to the experimental
data. This work will be reported in a separate paper.

\section{Conclusions}

We present a development of the self-consistent model of the Galactic CR halo, extending the model by \citet{dogiel2020}.
Our earlier model by \citet{dogiel2020} predicts a small size of the halo at low energies, which does not agree with
experimental data. To overcome this discrepancy, we include nonlinear Landau (NL) damping in the present model.

The key input parameters of the proposed model are the CR source $Q(p)$ as well as the spatial profiles of the vertical
magnetic field $B$, ionized gas density $n$, and temperature $T$. We show that all these parameters may significantly affect
the size of the halo, in particular at relatively low CR energies. The MHD-turbulence in the halo, which controls the
vertical escape of CRs, is entirely generated by the resonant CR-streaming instability. The equilibrium spectrum of MHD
waves in our present model is reached when the CR excitation rate is balanced by NL damping. This significantly suppresses
the amplitude of MHD waves compared to the model of \citet{dogiel2020}, thus making the halo size substantially larger.

We consider two alternative models of gas distributions in the halo: a single-component isothermal model and a two-phase
model composed of hot coronal gas and warm ionized gas. We showed that the single-component model requires very dense and
hot gas with $\beta \approx 40$ at low altitudes to be able to reproduce the experimental data. For the two-phase model, the
required gas parameters are much closer to those reported in the literature \citep[e.g.,][]{ferriere1998}.

Our model is able to reproduce the spectrum of CR protons in a wide range of energies, including the spectral features
observed between $\sim10$~GeV and $\sim10$~TeV (see Fig. \ref{fig:CR_spectra}). Despite some 20\% discrepancy with
experimental data below 10~GeV, our model predicts a reasonable halo size of about 1~kpc at 1~GeV. We argue that such a
discrepancy may be due to increasing influence of the Galactic disk at lower energies, which is still neglected in our
model.

\acknowledgments

The authors are grateful to an anonymous referee for constructive suggestions, and to Andrey Bykov for useful
discussions and comments. The work of DOC, VAD, ADE, and AMK is supported by the Russian Science Foundation via the Project
20-12-00047.

\section*{Note added in proof}
New data on CR proton spectrum reported by CALET \citep{adriani22} confirms the existence of the second spectral break at 10 TeV. The break position suggests that the scale height of hot gas should be about 2 kpc or less.

\bibliographystyle{apj}