# Neural Oblivious Decision Ensembles
for Deep Learning on Tabular Data

Sergei Popov
Yandex
<EMAIL_ADDRESS>
&Stanislav Morozov
Yandex
Lomonosov Moscow State University
<EMAIL_ADDRESS>
Artem Babenko
Yandex
National Research University
Higher School of Economics
<EMAIL_ADDRESS>
###### Abstract

Nowadays, deep neural networks (DNNs) have become the main instrument for
machine learning tasks within a wide range of domains, including vision, NLP,
and speech. Meanwhile, in an important case of heterogenous tabular data, the
advantage of DNNs over shallow counterparts remains questionable. In
particular, there is no sufficient evidence that deep learning machinery
allows constructing methods that outperform gradient boosting decision trees
(GBDT), which are often the top choice for tabular problems. In this paper, we
introduce Neural Oblivious Decision Ensembles (NODE), a new deep learning
architecture, designed to work with any tabular data. In a nutshell, the
proposed NODE architecture generalizes ensembles of oblivious decision trees,
but benefits from both end-to-end gradient-based optimization and the power of
multi-layer hierarchical representation learning. With an extensive
experimental comparison to the leading GBDT packages on a large number of
tabular datasets, we demonstrate the advantage of the proposed NODE
architecture, which outperforms the competitors on most of the tasks. We open-
source the PyTorch implementation of NODE and believe that it will become a
universal framework for machine learning on tabular data.

##  1 Introduction

The recent rise of deep neural networks (DNN) resulted in a substantial
breakthrough for a large number of machine learning tasks in computer vision,
natural language processing, speech recognition, reinforcement learning
(Goodfellow et al., 2016). Both gradient-based optimization via back-
propagation (Rumelhart et al., 1985) and hierarchical representation learning
appear to be crucial in increasing the performance of machine learning for
these problems by a large margin.

While the superiority of deep architectures in these domains is undoubtful,
machine learning for tabular data still did not fully benefit from the DNN
power. Namely, the state-of-the-art performance in problems with tabular
heterogeneous data is often achieved by ”shallow” models, such as gradient
boosted decision trees (GBDT) (Friedman, 2001; Chen & Guestrin, 2016; Ke et
al., 2017; Prokhorenkova et al., 2018). While the importance of deep learning
on tabular data is recognized by the ML community, and many works address this
problem (Zhou & Feng, 2017; Yang et al., 2018; Miller et al., 2017; Lay et
al., 2018; Feng et al., 2018; Ke et al., 2018), the proposed DNN approaches do
not consistently outperform the state-of-the-art shallow models by a notable
margin. In particular, to the best of our knowledge, there is still no
universal DNN approach that was shown to systematically outperform the leading
GBDT packages (e.g., XGBoost (Chen & Guestrin, 2016)). As additional evidence,
a large number of Kaggle ML competitions with tabular data are still won by
the shallow GBDT methods (Harasymiv, 2015). Overall, at the moment, there is
no dominant deep learning solution for tabular data problems, and we aim to
reduce this gap by our paper.

We introduce Neural Oblivious Decision Ensembles (NODE), a new DNN
architecture, designed to work with tabular problems. The NODE architecture is
partially inspired by the recent CatBoost package (Prokhorenkova et al.,
2018), which was shown to provide state-of-the-art performance on a large
number of tabular datasets. In a nutshell, CatBoost performs gradient boosting
on oblivious decision trees (decision tables) (Kohavi, 1994; Lou & Obukhov,
2017), which makes inference very efficient, and the method is quite resistant
to overfitting. In its essence, the proposed NODE architecture generalizes
CatBoost, making the splitting feature choice and decision tree routing
differentiable. As a result, the NODE architecture is fully differentiable and
could be incorporated in any computational graph of existing DL packages, such
as TensorFlow or PyTorch. Furthermore, NODE allows constructing multi-layer
architectures, which resembles ”deep” GBDT that is trained end-to-end, which
was never proposed before. Besides the usage of oblivious decision tables,
another important design choice is the recent entmax transformation (Peters et
al., 2019), which effectively performs a ”soft” splitting feature choice in
decision trees inside the NODE architecture. As discussed in the following
sections, these design choices are critical to obtain state-of-the-art
performance. In a large number of experiments, we compare the proposed
approach with the leading GBDT implementations with tuned hyperparameters and
demonstrate that NODE outperforms competitors consistently on most of the
datasets.

Overall, the main contributions of our paper can be summarized as follows:

  1. 1.

We introduce a new DNN architecture for machine learning on tabular data. To
the best of our knowledge, our method is the first successful example of deep
architectures that substantially outperforms leading GBDT packages on tabular
data.

  2. 2.

Via an extensive experimental evaluation on a large number of datasets, we
show that the proposed NODE architecture outperforms existing GBDT
implementations.

  3. 3.

The PyTorch implementation of NODE is available
online111https://github.com/Qwicen/node.

The rest of the paper is organized as follows. In Section 2 we review prior
work relevant to our method. The proposed Neural Oblivious Decision Ensembles
architecture is described in Section 3 and experimentally evaluated in Section
4. Section 5 concludes the paper.

##  2 Related work

In this section, we briefly review the main ideas from prior work that are
relevant to our method.

The state-of-the-art for tabular data. Ensembles of decision trees, such as
GBDT (Friedman, 2001) or random forests (Barandiaran, 1998), are currently the
top choice for tabular data problems. Currently, there are several leading
GBDT packages, such as XGBoost (Chen & Guestrin, 2016), LightGBM (Ke et al.,
2017), CatBoost (Prokhorenkova et al., 2018), which are widely used by both
academicians and ML practitioners. While these implementations vary in
details, on most of the tasks their performances do not differ much
(Prokhorenkova et al., 2018; Anghel et al., ). The most important distinction
of CatBoost is that it uses oblivious decision trees (ODTs) as weak learners.
As ODTs are also an important ingredient of our NODE architecture, we discuss
them below.

Oblivious Decision Trees. An oblivious decision tree is a regular tree of
depth $d$ that is constrained to use the same splitting feature and splitting
threshold in all internal nodes of the same depth. This constraint essentially
allows representing an ODT as a table with $2^{d}$ entries, corresponding to
all possible combinations of $d$ splits (Lou & Obukhov, 2017). Of course, due
to the constraints above, ODTs are significantly weaker learners compared to
unconstrained decision trees. However, when used in an ensemble, such trees
are less prone to overfitting, which was shown to synergize well with gradient
boosting (Prokhorenkova et al., 2018). Furthermore, the inference in ODTs is
very efficient: one can compute $d$ independent binary splits in parallel and
return the appropriate table entry. In contrast, non-oblivious decision trees
require evaluating $d$ splits sequentially.

Differentiable trees. The significant drawback of tree-based approaches is
that they usually do not allow end-to-end optimization and employ greedy,
local optimization procedures for tree construction. Thus, they cannot be used
as a component for pipelines, trained in an end-to-end fashion. To address
this issue, several works (Kontschieder et al., 2015; Yang et al., 2018; Lay
et al., 2018) propose to ”soften” decision functions in the internal tree
nodes to make the overall tree function and tree routing differentiable. In
our work, we advocate the usage of the recent entmax transformation (Peters et
al., 2019) to ”soften” decision trees. We confirm its advantages over the
previously proposed approaches in the experimental section.

Entmax. The key building block of our model is the entmax transformation
(Peters et al., 2019), which maps a vector of real-valued scores to a discrete
probability distribution. This transformation generalizes the traditional
softmax and its sparsity-enforcing alternative sparsemax (Martins & Astudillo,
2016), which has already received significant attention in a wide range of
applications: probabilistic inference, topic modeling, neural attention
(Niculae & Blondel, 2017; Niculae et al., 2018; Lin et al., 2019). The entmax
is capable to produce sparse probability distributions, where the majority of
probabilities are exactly equal to $0$. In this work, we argue that entmax is
also an appropriate inductive bias in our model, which allows differentiable
split decision construction in the internal tree nodes. Intuitively, entmax
can learn splitting decisions based on a small subset of data features (up to
one, as in classical decision trees), avoiding undesired influence from
others. As an additional advantage, using entmax for feature selection allows
for computationally efficient inference using the sparse pre-computed choice
vectors as described below in Section 3.

Multi-layer non-differentiable architectures. Another line of work (Miller et
al., 2017; Zhou & Feng, 2017; Feng et al., 2018) promotes the construction of
multi-layer architectures from non-differentiable blocks, such as random
forests or GBDT ensembles. For instance, (Zhou & Feng, 2017; Miller et al.,
2017) propose to use stacking of several random forests, which are trained
separately. In recent work, (Feng et al., 2018) introduces the multi-layer
GBDTs and proposes a training procedure that does not require each layer
component to be differentiable. While these works report marginal improvements
over shallow counterparts, they lack the capability for end-to-end training,
which could result in inferior performance. In contrast, we argue that end-to-
end training is crucial and confirm this claim in the experimental section.

Specific DNN for tabular data. While a number of prior works propose
architectures designed for tabular data (Ke et al., 2018; Shavitt & Segal,
2018), they mostly do not compare with the properly tuned GBDT
implementations, which are the most appropriate baselines. The recent preprint
(Ke et al., 2018) reports the marginal improvement over GBDT with default
parameters, but in our experiments, the baseline performance is much higher.
To the best of our knowledge, our approach is the first to consistently
outperform the tuned GBDTs over a large number of datasets.

##  3 Neural Oblivious Decision Ensembles

We introduce the Neural Oblivious Decision Ensemble (NODE) architecture with a
layer-wise structure similar to existing deep learning models. In a nutshell,
our architecture consists of differentiable oblivious decision trees (ODT)
that are trained end-to-end by backpropagation. We describe our implementation
of the differentiable NODE layer in Section 3.1, the full model architecture
in Section 3.2, and the training and inference procedures in section 3.3.

###  3.1 Differentiable Oblivious Decision Trees

The core building block of our model is a Neural Oblivious Decision Ensemble
(NODE) layer. The layer is composed of $m$ differentiable oblivious decision
trees (ODTs) of equal depth $d$. As an input, all $m$ trees get a common
vector $x\in\mathbb{R}^{n}$, containing $n$ numeric features. Below we
describe a design of a single differentiable ODT.

Figure 1: The single ODT inside the NODE layer. The splitting features and the
splitting thresholds are shared across all the internal nodes of the same
depth. The output is a sum of leaf responses scaled by the choice weights.

In its essence, an ODT is a “decision table” that splits the data along $d$
splitting features and compares each feature to a learned threshold. Then, the
tree returns one of the $2^{d}$ possible responses, corresponding to the
comparisons result. Therefore, each ODT is completely determined by its
splitting features $f\in\mathbb{R}^{d}$, splitting thresholds
$b\in\mathbb{R}^{d}$ and a $d$-dimensional tensor of responses
$R\in\mathbb{R}^{{\scaleobj{.7}{\displaystyle\underbrace{2\times
2\times…2}_{d}}}}$. In this notation, the tree output is defined as:

$h(x)=R[\mathbbm{1}(f_{1}(x)-b_{1}),\dots,\mathbbm{1}(f_{d}(x)-b_{d})],$ (1)

where $\mathbbm{1}(\cdot)$ denotes the Heaviside function.

To make the tree output (1) differentiable, we replace the splitting feature
choice $f_{i}$ and the comparison operator $\mathbbm{1}(f_{i}(x)-b_{i})$ by
their continuous counterparts. There are several existing approaches that can
be used for modelling differentiable choice functions in decision trees (Yang
et al., 2018), for instance, REINFORCE (Williams, 1992) or Gumbel-softmax
(Jang et al., 2016). However, these approaches typically require long training
time, which can be crucial in practice.

Instead, we propose to use the $\alpha$-entmax transformation (Peters et al.,
2019) as it is able to learn sparse choices, depending only on a few features,
via standard gradient descent. The choice function is hence replaced by a
weighted sum of features, with weights computed as entmax over the learnable
feature selection matrix $F\in\mathbb{R}^{d\times n}$:

$\hat{f}_{i}(x)=\sum_{j=1}^{n}x_{j}\cdot entmax_{\alpha}(F_{ij})\vspace{-2mm}$
(2)

Similarly, we relax the Heaviside function $\mathbbm{1}(f_{i}(x)-b_{i})$ as a
two-class entmax, which we denote as
$\sigma_{\alpha}(x){=}entmax_{\alpha}(\left[x,0\right])$. As different
features can have different characteristic scales, we use the scaled version
$c_{i}(x)=\sigma_{\alpha}\left(\frac{f_{i}(x)-b_{i}}{\tau_{i}}\right)$, where
$b_{i}$ and $\tau_{i}$ are learnable parameters for thresholds and scales
respectively.

Based on the $c_{i}(x)$ values, we define a ”choice” tensor
$C\in\mathbb{R}^{{{\scaleobj{.7}{\displaystyle\underbrace{2\times
2\times…2}_{d}}}}}$ of the same size as the response tensor $R$ by computing
the outer product of all $c_{i}$:

$C(x)=\begin{bmatrix}c_{1}(x)\\\
1-c_{1}(x)\end{bmatrix}\otimes\begin{bmatrix}c_{2}(x)\\\
1-c_{2}(x)\end{bmatrix}\otimes\dots\otimes\begin{bmatrix}c_{d}(x)\\\
1-c_{d}(x)\end{bmatrix}$ (3)

The final prediction is then computed as a weighted linear combination of
response tensor entries $R$ with weights from the entries of choice tensor
$C$:

$\hat{h}(x)=\sum_{i_{1},\dots i_{d}\in\\{0,1\\}^{d}}R_{i_{1},\dots,i_{d}}\cdot
C_{i_{1},\dots,i_{d}}(x)\vspace{-1mm}$ (4)

Note, that this relaxation equals to the classic non-differentiable ODT
$h(x)$(1) iff both feature selection and threshold functions reach “one-hot”
state, i.e. entmax always returns non-zero weights for a single feature and
$c_{i}$ always return exactly zeros or ones.

Finally, the output of the NODE layer is composed as a concatenation of the
outputs of $m$ individual trees
$\left[\hat{h}_{1}(x),\dots,\hat{h}_{m}(x)\right]$.

Multidimensional tree outputs. In the description above, we assumed that tree
outputs are one-dimensional $\hat{h}(x)\in\mathbb{R}$. For classification
problems, where NODE predicts probabilities of each class, we use
multidimensional tree outputs $\hat{h}(x)\in\mathbb{R}^{|C|}$, where $|C|$ is
a number of classes.

###  3.2 Going deeper with the NODE architecture

The NODE layer, described above, can be trained alone or within a complex
structure, like fully-connected layers that can be organized into the multi-
layer architectures. In this work, we introduce a new architecture, following
the popular DenseNet (Huang et al., 2017) model and train it end-to-end via
backpropagation.

Figure 2: The NODE architecture, consisting of densely connected NODE layers.
Each layer contains several trees whose outputs are concatenated and serve as
input for the subsequent layer. The final prediction is obtained by averaging
the outputs of all trees from all the layers.

Similar to DenseNet, our architecture is a sequence of $k$ NODE layers (see
Section 3.1), where each layer uses a concatenation of all previous layers as
its input. The input “layer 0” of this architecture corresponds to the input
features $x$, accessible by all successor layers. Due to such a design, our
architecture is capable to learn both shallow and deep decision rules. A
single tree on $i$-th layer can rely on chains of up to $i-1$ layer outputs as
features, allowing it to capture complex dependencies. The resulting
prediction is a simple average of all decision trees from all layers.

Note, in the multi-layer architecture described above, tree outputs
$\hat{h}(x)$ from early layers are used as inputs for subsequent layers.
Therefore, we do not restrict the dimensionality of $\hat{h}(x)$ to be equal
to the number of classes, and allow it to have an arbitrary dimensionality
$l$, which correspond to the $(d+1)$-dimensional response tensor
$R\in\mathbb{R}^{{\scaleobj{.7}{\displaystyle\underbrace{2\times
2\times…2}_{d}\times l}}}$. When averaging the predictions from all layers,
only first $|C|$ coordinates of $\hat{h}(x)$ are used for classification
problems and the first one for regression problems. Overall, $l$ is an
additional hyperparameter with typical values in $\left[1,3\right]$.

###  3.3 Training

Here we summarize the details of our training protocol.

Data preprocessing. First, we transform each data feature to follow a normal
distribution via quantile
transform222sklearn.preprocessing.QuantileTransformer. In experiments, we
observed that this step was important for stable training and faster
convergence.

Initialization. Before training, we perform the data-aware initialization
(Mishkin & Matas, 2016) to obtain a good initial parameter values. In
particular, we initialize the feature selection matrix uniformly $F_{ij}\sim
U(0,1)$, while the thresholds $b$ are initialized with random feature values
$f_{i}(x)$ observed in the first data batch. The scales $\tau_{i}$ are
initialized in such a way that all the samples in the first batch belong to
the linear region of $\sigma_{\alpha}$, and hence receive nonzero gradients.
Finally, the response tensor entries are initialized with the standard normal
distribution $R[i_{1},\dots,i_{d}]\sim N(0,1)$.

Training. As for existing DNN architectures, NODE is trained end-to-end via
mini-batch SGD. We jointly optimize all model parameters: $F,b,R$. In this
work, we experimented with traditional objective functions (cross-entropy for
classification and mean squared error for regression), but any differentiable
objective can be used as well. As an optimization method, we use the recent
Quasi-Hyperbolic Adam with parameters recommended in the original paper (Ma &
Yarats, 2018). We also average the model parameters over $c=5$ consecutive
checkpoints (Izmailov et al., 2018) and pick the optimal stopping point on the
hold-out validation dataset.

Inference. During training, a significant fraction of time is spent computing
the entmax function and multiplying the choice tensor. Once the model is
trained, one can pre-compute entmax feature selectors and store them as a
sparse vector (e.g., in coordinate (coo) format), making inference more
efficient.

##  4 Experiments

In this section, we report the results of a comparison between our approach
and the leading GBDT packages. We also provide several ablation studies that
demonstrate the influence of each design choice in the proposed NODE
architecture.

###  4.1 Comparison to the state-of-the-art.

As our main experiments, we compare the proposed NODE architecture with two
state-of-the-art GBDT implementations on a large number of datasets. In all
the experiments we set $\alpha$ parameter in the entmax transformation to
$1.5$. All other details of the comparison protocol are described below.

Datasets. We perform most of the experiments on six open-source tabular
datasets from different domains: Epsilon, YearPrediction, Higgs, Microsoft,
Yahoo, Click. The detailed description of the datasets is available in
appendix. All the datasets provide train/test splits, and we used 20% samples
from the train set as a validation set to tune the hyperparameters. For each
dataset, we fix the train/val/test splits for a fair comparison. For the
classification datasets (Epsilon, Higgs, Click), we minimize cross-entropy
loss and report the classification error. For the regression and ranking
datasets (YearPrediction, Microsoft, Yahoo), we minimize and report mean
squared error (which corresponds to the pointwise approach to learning-to-
rank).

Methods. We compare the proposed NODE architecture to the following baselines:

  * •

Catboost. The recent GBDT implementation (Prokhorenkova et al., 2018) that
uses oblivious decision trees as weak learners. We use the open-source
implementation, provided by the authors.

  * •

XGBoost. The most popular GBDT implementation widely used in machine learning
competitions (Chen & Guestrin, 2016). We use the open-source implementation,
provided by the authors.

  * •

FCNN. Deep neural network, consisting of several fully-connected layers with
ReLU non-linearity layers (Nair & Hinton, 2010).

Regimes. We perform comparison in two following regimes that are the most
important in practice:

  * •

Default hyperparameters. In this regime, we compare the methods as easy-to-
tune toolkits that could be used by a non-professional audience. Namely, here
we do not tune hyperparameters and use the default ones provided by the GBDT
packages. The only tunable parameter here is a number of trees (up to 2048) in
CatBoost/XGBoost, which is set based on the validation set. We do not compare
with FCNN in this regime, as it typically requires much tuning, and we did not
find the set of parameters, appropriate for all datasets. The default
architecture in our model contains only a single layer with 2048 decision
trees of depth six. Both of these hyperparameters were inherited from the
CatBoost package settings for oblivious decision trees. With these parameters,
the NODE architecture is shallow, but it still benefits from end-to-end
training via back-propagation.

  * •

Tuned hyperparameters. In this regime, we tune the hyperparameters for both
NODE and the competitors on the validation subsets. The optimal configuration
for NODE contains between two and eight NODE layers, while the total number of
trees across all the layers does not exceed 2048. The details of
hyperparameter optimization are provided in appendix.

The results of the comparison are summarized in Table 1 and Table 2. For all
methods, we report mean performance and standard deviations computed over ten
runs with different random seeds. Several key observations are highlighted
below:

| Epsilon | YearPrediction | Higgs | Microsoft | Yahoo | Click
---|---|---|---|---|---|---
Default hyperparameters
CatBoost | $0.1119{\pm}{2e{-}4}$ | $80.68{\pm}0.04$ | $0.2434{\pm}{2e{-}4}$ | $0.5587{\pm}{2e{-}4}$ | $0.5781{\pm}{3e{-}4}$ | $0.3438{\pm}{1e{-}3}$
XGBoost | $0.1144$ | $81.11$ | $0.2600$ | $0.5637$ | $0.5756$ | $0.3461$
NODE | $0.1043{\pm}{4e{-}4}$ | $77.43{\pm}0.09$ | $0.2412{\pm}{5e{-}4}$ | $0.5584{\pm}{3e{-}4}$ | $0.5666{\pm}{5e{-}4}$ | $0.3309{\pm}{3e{-}4}$

Table 1: The comparison of NODE with the shallow state-of-the-art counterparts
with default hyperparameters. The results are computed over ten runs with
different random seeds.

| Epsilon | YearPrediction | Higgs | Microsoft | Yahoo | Click
---|---|---|---|---|---|---
Tuned hyperparameters
CatBoost | $0.1113{\pm}{4e{-}4}$ | $79.67{\pm}0.12$ | $0.2378{\pm}{1e{-}4}$ | $0.5565{\pm}{2e{-}4}$ | $0.5632{\pm}{3e{-}4}$ | $0.3401{\pm}{2e{-}3}$
XGBoost | $0.1112{\pm}{6e{-}4}$ | $78.53{\pm}0.09$ | $0.2328{\pm}{3e{-}4}$ | $0.5544{\pm}{1e{-}4}$ | $0.5420{\pm}{4e{-}4}$ | $0.3334{\pm}{2e{-}3}$
FCNN | $0.1041{\pm}{2e{-}4}$ | $79.99{\pm}0.47$ | $0.2140{\pm}{2e{-}4}$ | $0.5608{\pm}{4e{-}4}$ | $0.5773{\pm}{1e{-}3}$ | $0.3325{\pm}{2e{-}3}$
NODE | $0.1034{\pm}{3e{-}4}$ | $76.21{\pm}0.12$ | $0.2101{\pm}{5e{-}4}$ | $0.5570{\pm}{2e{-}4}$ | $0.5692{\pm}{2e{-}4}$ | $0.3312{\pm}{2e{-}3}$
mGBDT | OOM | $80.67$ | OOM | OOM | OOM | OOM
DeepForest | $0.1179$ | — | $0.2391$ | — | — | $0.3333$

Table 2: The comparison of NODE with both shallow and deep counterparts with
hyperparameters tuned for optimal performance. The results are computed over
ten runs with different random seeds.

  1. 1.

With default hyperparameters, the proposed NODE architecture consistently
outperforms both CatBoost and XGBoost on all datasets. The results advocate
the usage of NODE as a handy tool for machine learning on tabular problems.

  2. 2.

With tuned hyperparameters, NODE also outperforms the competitors on most of
the tasks. Two exceptions are the Yahoo and Microsoft datasets, where tuned
XGBoost provides the highest performance. Given the large advantage of XGBoost
over CatBoost on Yahoo, we speculate that the usage of oblivious decision
trees is an inappropriate inductive bias for this dataset. This implies that
NODE should be extended to non-oblivious trees, which we leave for future
work.

  3. 3.

In the regime with tuned hyperparameters on some datasets FCNN outperforms
GBDT, while on others GBDT is superior. Meanwhile, the proposed NODE
architecture appears to be a universal instrument, providing the highest
performance on most of the tasks.

For completeness we also aimed to compare to previously proposed architectures
for deep learning on tabular data. Unfortunately, many works did not publish
the source code. We were only able to perform a partial comparison with mGBDT
(Feng et al., 2018) and DeepForest (Zhou & Feng, 2017), which source code is
available. For both baselines, we use the implementations, provided by the
authors, and tune the parameters on the validation set. Note, that the
DeepForest implementation is available only for classification problems.
Moreover, both implementations do not scale well, and for many datasets, we
obtained Out-Of-Memory error (OOM). On datasets in our experiments it turns
out that properly tuned GBDTs outperform both (Feng et al., 2018) and (Zhou &
Feng, 2017).

###  4.2 Ablative analysis

In this section, we analyze the key architecture components that define our
model.

Choice functions. Constructing differentiable decision trees requires a
function that selects items from a set. Such function is required for both
splitting feature selection and decision tree routing. We experimented with
four possible options, each having different implications:

  * •

Softmax learns “dense” decision rules where all items have nonzero weights;

  * •

Gumbel-Softmax (Jang et al., 2016) learns to stochastically sample a single
element from a set;

  * •

Sparsemax (Martins & Astudillo, 2016) learns sparse decision rules, where only
a few items have nonzero weights;

  * •

Entmax (Peters et al., 2019) generalizes both sparsemax and softmax; it is
able to learn sparse decision rules, but is smoother than sparsemax, being
more appropriate for gradient-based optimization. In comparison $\alpha$
parameter was set to $1.5$.

We experimentally compare the four options above with both shallow and deep
architectures in Table 3. We use the same choice function for both feature
selection and tree routing across all experiments. In Gumbel-Softmax, we
replaced it with hard argmax one-hot during inference. The results clearly
show that Entmax with $\alpha{=}1.5$ outperforms the competitors across all
experiments. First, Table 3 demonstrates that sparsemax and softmax are not
universal choice functions. For instance, on the YearPrediction dataset,
sparsemax outperforms softmax, while on the Epsilon dataset softmax is
superior. In turn, entmax provides great empirical performance across all
datasets. Another observation is that Gumbel-Softmax is unable to learn deep
architectures with both constant and annealed temperature schedules. This
behavior is probably caused by the stochasticity of Gumbel-Softmax and the
responses on the former layers are too noisy to produce useful features for
the latter layers.

Dataset | YearPrediction | Epsilon
---|---|---
Function | softmax | Gumbel | sparsemax | entmax | softmax | Gumbel | sparsemax | entmax
1 layer | 78.41 | 79.39 | 78.13 | 77.43 | 0.1045 | 0.1979 | 0.1083 | 0.1043
2 layers | 77.61 | 79.31 | 76.81 | 77.05 | 0.1041 | 0.2884 | 0.1052 | 0.1031
4 layers | 77.58 | 79.69 | 76.60 | 76.21 | 0.1034 | 0.2908 | 0.1058 | 0.1033
8 layers | 77.47 | 80.49 | 76.31 | 76.17 | 0.1036 | 0.3081 | 0.1058 | 0.1036

Table 3: The experimental comparison of various choice functions and
architecture depths. The values represent mean squared error for
YearPrediction and classification error rate for Epsilon.

Feature importance. In this series of experiments, we analyze the internal
representations, learned by the NODE architecture. We begin by estimating the
feature importances from different layers of a multi-layer ensemble via
permutation feature importance, initially introduced in (Breiman, 2001).
Namely, for 10.000 objects from the Higgs dataset we randomly shuffle the
values of each feature (original or learnt on some NODE layer) and compute the
increase in the classification error. Then for each layer, we split feature
importance values into seven equal bins and calculate the total feature
importance of each bin, shown on Figure 3 (left-top). We discovered that the
features from the first layer are used the most, with feature importances
decreasing with depth. This figure shows that deep layers are able to produce
important features, even though earlier layers have an advantage because of
the DenseNet architecture. Next, we estimated the mean absolute contribution
of individual trees to the final response, reported on Figure 3 (left-bottom).
One can see the reverse trend, deep trees tend to contribute more to the final
response. Figure 3 (right) clearly shows that there is anticorrelation of
feature importances and contributions in the final response, which implies
that the main role of ealier layers is to produce informative features, while
the latter layers mostly use them for accurate prediction.

|
---|---
Figure 3: NODE on UCI Higgs dataset: Left-Top: individual feature importance
distributions for both original and learned features. Left-Bottom: mean
absolute contribution of individual trees to the final response. Right:
responses dependence on feature importances. See details in the text.

Training/Inference runtime. Finally, we compare the NODE runtime to the
timings of the state-of-the-art GBDT implementations. In Table 4 we report the
training and inference time for million of objects from the YearPrediction
dataset. In this experiment, we evaluate ensembles of 1024 trees of depth six
with all other parameters set to their default values. Our GPU setup has a
single 1080Ti GPU and 2 CPU cores. In turn, our CPU setup has a 28-core Xeon
E5-2660 v4 processor (which costs almost twice as much as the GPU). We use
CatBoost v0.15 and XGBoost v0.90 as baselines, while NODE inference runs on
PyTorch v1.1.0. Overall, NODE inference time is on par with heavily optimized
GBDT libraries despite being implemented in pure PyTorch (i.e. no custom
kernels).

Method | NODE 8 layers 1080Ti | XGBoost Xeon | XGBoost 1080Ti | CatBoost Xeon
---|---|---|---|---
Training | 7min 42s | 5min 39s | 1min 13s | 41s
Inference | 8.56s | 5.94s | 4.45s | 4.62s
Table 4: Training and inference runtime for models with 1024 trees of depth
six on the YearPrediction dataset, averaged over five runs. Both timings of
eight-layer NODE architecture are on par with timings of shallow counterparts
of the same total number of trees in an ensemble.

##  5 Conclusion

In this paper, we introduce a new DNN architecture for deep learning on
heterogeneous tabular data. The architecture is differentiable deep GBDTs,
trained end-to-end via backpropagation. In extensive experiments, we
demonstrate the advantages of our architecture over existing competitors with
the default and tuned hyperparameters. A promising research direction is
incorporating the NODE layer into complex pipelines trained via back-
propagation. For instance, in multi-modal problems, the NODE layer could be
employed as a way to incorporate the tabular data, as CNNs are currently used
for images, or RNNs are used for sequences.

## References

  * (1) Andreea Anghel, Nikolaos Papandreou, Thomas Parnell, Alessandro de Palma, and Haralampos Pozidis.  Benchmarking and optimization of gradient boosting decision tree algorithms.
  * Barandiaran (1998) Iñigo Barandiaran.  The random subspace method for constructing decision forests.  _IEEE transactions on pattern analysis and machine intelligence_ , 1998.
  * Breiman (2001) Leo Breiman.  Random forests.  _Machine Learning_ , 45(1):5–32, 2001.  doi: 10.1023/A:1010933404324.  URL https://doi.org/10.1023/A:1010933404324.
  * Chen & Guestrin (2016) Tianqi Chen and Carlos Guestrin.  Xgboost: A scalable tree boosting system.  In _Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining_ , pp. 785–794, 2016.
  * Feng et al. (2018) Ji Feng, Yang Yu, and Zhi-Hua Zhou.  Multi-layered gradient boosting decision trees.  In _Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS_ , 2018.
  * Friedman (2001) Jerome H Friedman.  Greedy function approximation: a gradient boosting machine.  _Annals of statistics_ , pp. 1189–1232, 2001.
  * Goodfellow et al. (2016) Ian Goodfellow, Yoshua Bengio, and Aaron Courville.  _Deep learning_.  MIT press, 2016.
  * Harasymiv (2015) Vasyl Harasymiv.  Lessons from 2 million machine learning models on kaggle, 2015.
  * Huang et al. (2017) Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2017.
  * Izmailov et al. (2018) Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.  Averaging weights leads to wider optima and better generalization.  _arXiv preprint arXiv:1803.05407_ , 2018.
  * Jang et al. (2016) Eric Jang, Shixiang Gu, and Ben Poole.  Categorical reparameterization with gumbel-softmax.  _CoRR_ , abs/1611.01144, 2016.  URL http://dblp.uni-trier.de/db/journals/corr/corr1611.html#JangGP16.
  * Ke et al. (2017) Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, Wei Chen, Weidong Ma, Qiwei Ye, and Tie-Yan Liu.  Lightgbm: A highly efficient gradient boosting decision tree.  In _Advances in Neural Information Processing Systems_ , pp. 3146–3154, 2017.
  * Ke et al. (2018) Guolin Ke, Jia Zhang, Zhenhui Xu, Jiang Bian, and Tie-Yan Liu.  Tabnn: A universal neural network solution for tabular data.  2018\.
  * Kohavi (1994) Ron Kohavi.  Bottom-up induction of oblivious read-once decision graphs: strengths and limitations.  In _AAAI_ , 1994.
  * Kontschieder et al. (2015) Peter Kontschieder, Madalina Fiterau, Antonio Criminisi, and Samuel Rota Bulo.  Deep neural decision forests.  In _Proceedings of the IEEE international conference on computer vision_ , 2015.
  * Lay et al. (2018) Nathan Lay, Adam P Harrison, Sharon Schreiber, Gitesh Dawer, and Adrian Barbu.  Random hinge forest for differentiable learning.  _arXiv preprint arXiv:1802.03882_ , 2018.
  * Lin et al. (2019) Tianyi Lin, Zhiyue Hu, and Xin Guo.  Sparsemax and relaxed wasserstein for topic sparsity.  In _Proceedings of the Twelfth ACM International Conference on Web Search and Data Mining_ , 2019.
  * Lou & Obukhov (2017) Yin Lou and Mikhail Obukhov.  Bdt: Gradient boosted decision tables for high accuracy and scoring efficiency.  In _Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , 2017.
  * Ma & Yarats (2018) Jerry Ma and Denis Yarats.  Quasi-hyperbolic momentum and adam for deep learning.  _arXiv preprint arXiv:1810.06801_ , 2018.
  * Martins & Astudillo (2016) Andre Martins and Ramon Astudillo.  From softmax to sparsemax: A sparse model of attention and multi-label classification.  In _International Conference on Machine Learning_ , 2016.
  * Miller et al. (2017) Kevin Miller, Chris Hettinger, Jeffrey Humpherys, Tyler Jarvis, and David Kartchner.  Forward thinking: building deep random forests.  _arXiv preprint arXiv:1705.07366_ , 2017.
  * Mishkin & Matas (2016) Dmytro Mishkin and Jiri Matas.  All you need is a good init.  In _4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings_ , 2016.
  * Nair & Hinton (2010) Vinod Nair and Geoffrey E Hinton.  Rectified linear units improve restricted boltzmann machines.  In _Proceedings of the 27th international conference on machine learning (ICML-10)_ , 2010.
  * Niculae & Blondel (2017) Vlad Niculae and Mathieu Blondel.  A regularized framework for sparse and structured neural attention.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Niculae et al. (2018) Vlad Niculae, André FT Martins, Mathieu Blondel, and Claire Cardie.  Sparsemap: Differentiable sparse structured inference.  _arXiv preprint arXiv:1802.04223_ , 2018.
  * Peters et al. (2019) Ben Peters, Vlad Niculae, and André F. T. Martins.  Sparse sequence-to-sequence models.  In _ACL, 2019_ , pp. 1504–1519, 2019.
  * Prokhorenkova et al. (2018) Liudmila Prokhorenkova, Gleb Gusev, Aleksandr Vorobev, Anna Veronika Dorogush, and Andrey Gulin.  Catboost: unbiased boosting with categorical features.  In _Advances in Neural Information Processing Systems_ , pp. 6638–6648, 2018.
  * Rumelhart et al. (1985) David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams.  Learning internal representations by error propagation.  Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.
  * Shavitt & Segal (2018) Ira Shavitt and Eran Segal.  Regularization learning networks: Deep learning for tabular datasets.  In _Advances in Neural Information Processing Systems_ , pp. 1379–1389, 2018.
  * Williams (1992) Ronald J. Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine Learning_ , 8:229–256, 1992.
  * Yang et al. (2018) Yongxin Yang, Irene Garcia Morillo, and Timothy M Hospedales.  Deep neural decision trees.  _arXiv preprint arXiv:1806.06988_ , 2018.
  * Zhou & Feng (2017) Zhi-Hua Zhou and Ji Feng.  Deep forest: Towards an alternative to deep neural networks.  In _Proceedings of the Twenty-Sixth International Joint Conference on Artificial Intelligence, IJCAI_ , 2017.

##  Appendix A Appendix

###  A.1 Description of the datasets

In our experiments, we used six tabular datasets, described in Table 5. (1)
Epsilon is high dimensional dataset from the PASCAL Large Scale Learning
Challenge 2008. The problem is a binary classification. (2) YearPrediction is
a subset of Million Song Dataset. It is regression dataset, and the task is to
predict the release year of the song by using the audio features. It contains
tracks from 1922 to 2011. (3) Higgs is a dataset from the UCI ML Repository.
The problem is to predict whether the given event produces Higgs bosons or
not. (4) Microsoft is a Learning to Rank Dataset. It consists of
136-dimensional feature vectors extracted from query-url pairs. Each pair has
relevance judgment labels, which take values from 0 (irrelevant) to 4
(perfectly relevant) (5) Yahoo is very similar ranking dataset with query-url
pairs labeled from 0 to 4. We treat both ranking problems as regression (which
corresponds to the pointwise approach to learning-to-rank) (6) Click is a
subset of data from the 2012 KDD Cup. For the subset construction, we randomly
sample 500.000 objects of a positive class and 500.000 objects of a negative
class. The categorical features were converted to numerical ones via Leave-
One-Out encoder from category_encoders package of the scikit-learn library.

###  A.2 Optimization of hyperparameters

In order to tune the hyperparameters, we performed a random stratified split
of full training data into train set (80%) and validation set (20%) for the
Epsilon, YearPrediction, Higgs, Microsoft, and Click datasets. For Yahoo, we
use train/val/test split provided by the dataset authors. We use the
Hyperopt333https://github.com/hyperopt/hyperopt library to optimize Catboost,
XGBoost, and FCNN hyperparameters. For each method, we perform 50 steps of
Tree-structured Parzen Estimator (TPE) optimization algorithm. As a final
configuration, we choose the set of hyperparameters, corresponding to the
smallest loss on the validation set.

####  A.2.1 Catboost and XGBoost

On each iteration of Hyperopt, the number of trees was set based on the
validation set, with maximal trees count set to 2048. Below is the list of
hyperparameters and their search spaces for Catboost.

  * •

learning_rate: Log-Uniform distribution $[e^{-5},1]$

  * •

random_strength: Discrete uniform distribution $[1,20]$

  * •

one_hot_max_size: Discrete uniform distribution $[0,25]$

  * •

l2_leaf_reg: Log-Uniform distribution $[1,10]$

  * •

bagging_temperature: Uniform distribution $[0,1]$

  * •

leaf_estimation_iterations: Discrete uniform distribution $[1,10]$

| Train | Test | Features | Task | Metric | Description
---|---|---|---|---|---|---
Epsilon444https://www.csie.ntu.edu.tw/ cjlin/libsvmtools/datasets/binary.html | 400K | 100K | 2000 | Classification | Error | PASCAL Challenge 2008
YearPrediction555https://archive.ics.uci.edu/ml/datasets/yearpredictionmsd | 463K | 51.6K | 90 | Regression | MSE | Million Song Dataset
Higgs666https://archive.ics.uci.edu/ml/datasets/HIGGS | 10.5M | 500K | 28 | Classification | Error | UCI ML Higgs
Microsoft777https://www.microsoft.com/en-us/research/project/mslr/ | 723K | 241K | 136 | Regression | MSE | MSLR-WEB10K
Yahoo888https://webscope.sandbox.yahoo.com/catalog.php?datatype=c | 544K | 165K | 699 | Regression | MSE | Yahoo LETOR dataset
Click999http://www.kdd.org/kdd-cup/view/kdd-cup-2012-track-2 | 800K | 200K | 11 | Classification | Error | 2012 KDD Cup
Table 5: The datasets used in our experiments.

XGBoost tuned parameters and their search spaces:

  * •

eta: Log-Uniform distribution $[e^{-7},1]$

  * •

max_depth: Discrete uniform distribution $[2,10]$

  * •

subsample: Uniform distribution $[0.5,1]$

  * •

colsample_bytree: Uniform distribution $[0.5,1]$

  * •

colsample_bylevel: Uniform distribution $[0.5,1]$

  * •

min_child_weight: Log-Uniform distribution $[e^{-16},e^{5}]$

  * •

alpha: Uniform choice {0, Log-Uniform distribution $[e^{-16},e^{2}]$}

  * •

lambda: Uniform choice {0, Log-Uniform distribution $[e^{-16},e^{2}]$}

  * •

gamma: Uniform choice {0, Log-Uniform distribution $[e^{-16},e^{2}]$}

####  A.2.2 FCNN

Fully connected neural networks were tuned using Hyperas
101010https://github.com/maxpumperla/hyperas library, which is a Keras wrapper
for Hyperopt. We consider FCNN constructed from the following blocks: Dense-
ReLU-Dropout. The number of units in each layer is independent of each other,
and dropout value is the same for the whole network. The networks are trained
with the Adam optimizer with averaging the model parameters over $c{=}5$
consecutive checkpoints (Izmailov et al., 2018) and early stopping on
validation. Batch size is fixed to 1024 for all datasets. Below is the list of
tuned hyperparameters.

  * •

Number of layers: Discrete uniform distribution $[2,7]$

  * •

Number of units: Dicrete uniform distribution over a set
$\\{128,256,512,1024\\}$

  * •

Learning rate: Uniform distribution $[1e-4,1e-2]$

  * •

Dropout: Uniform distribution $[0,0.5]$

####  A.2.3 NODE

Neural Oblivious Decision Ensembles were tuned by grid search over the
following hyperparameter values. In the multi-layer NODE, we use the same
architecture for all layers, i.e., the same number of trees of the same depth.
total_tree_count here denotes the total number of trees on all layers. For
each dataset, we use the maximal batch size, which fits in the GPU memory. We
always use learning rate $10^{-3}$.

  * •

num_layers: {2, 4, 8}

  * •

total_tree_count: {1024, 2048}

  * •

tree_depth: {6, 8}

  * •

tree_output_dim: {2, 3}


# AdaptSim: Task-Driven Simulation Adaptation
for Sim-to-Real Transfer

Allen Z. Ren1 Hongkai Dai2 Benjamin Burchfiel2 Anirudha Majumdar1  1 Princeton
University 2 Toyota Research Institute

###### Abstract

Simulation parameter settings such as contact models and object geometry
approximations are critical to training robust robotic policies capable of
transferring from simulation to real-world deployment. Previous approaches
typically handcraft distributions over such parameters (domain randomization),
or identify parameters that best match the dynamics of the real environment
(system identification). However, there is often an _irreducible gap_ between
simulation and reality: attempting to match the dynamics between simulation
and reality across all states and tasks may be infeasible and may not lead to
policies that perform well in reality for a specific task. Addressing this
issue, we propose AdaptSim, a new _task-driven_ adaptation framework for sim-
to-real transfer that aims to optimize task performance in target (real)
environments — instead of matching dynamics between simulation and reality.
First, we meta-learn an adaptation policy in simulation using reinforcement
learning for adjusting the simulation parameter distribution based on the
current policy’s performance in a target environment. We then perform
iterative real-world adaptation by inferring new simulation parameter
distributions for policy training, using a small amount of real data. We
perform experiments in three robotic tasks: (1) swing-up of linearized double
pendulum, (2) dynamic table-top pushing of a bottle, and (3) dynamic scooping
of food pieces with a spatula. Our extensive simulation and hardware
experiments demonstrate AdaptSim achieving 1-3x asymptotic performance and
$\sim$2x real data efficiency when adapting to different environments,
compared to methods based on Sys-ID and directly training the task policy in
target environments.111video available at https://youtu.be/p2msMCOFDDg

##  I Introduction

Learning robust policies for real-world manipulation tasks typically requires
a substantial amount of training data. Since using real data exclusively can
be very expensive or even infeasible, we often resort to training mostly in
simulation. This raises the question: how should we specify simulation
parameters to maximize performance in the real world while minimizing the
amount of real-world data we require?

A popular method among the robotics community is to perform _domain
randomization_ : train a policy using a wide range of different simulation
parameters in the hope that the policy can thus handle possible real-world
variations in dynamics or observations. This approach has enabled some of the
most prominent successes in sim-to-real transfer of robotic policies
[akkaya2019solving, handa2022dextreme, lee2020learning, loquercio2019deep].
However, the trained policy may achieve good _average_ performance, but
perform poorly in a particular real environment. Addressing this issue, there
has been work in performing system identification (Sys-ID) for providing a
point or a distributional estimate of parameters that best matches the robot
or environment dynamics exhibited in real-world data. This estimation can be
performed using either a single iteration [lim2021planar] or multiple ones
[ramos2019bayessim]. These _adaptive_ domain randomization techniques allow
training policies suited to specific target environments.

Figure 1: AdaptSim iteratively infers a better simulation parameter
distribution for training the task policy in simulation based on real
trajectory observations of executing a task policy. When dynamically scooping
up chocolate raisins, the updated task policy learns to use smaller initial
pitch angle of the spatula from the table, such that the pieces do not bounce
off after contact.

While simple objects such as a box and its properties like the inertia can be
well-modeled, there is a substantial amount of “irreducible” sim-to-real gap
in many settings such as contact-rich manipulation tasks. Consider the task of
using a cooking spatula to dynamically scoop up small pieces of food from a
table (Fig. 1). The exact object geometry of the pieces and spatula is
difficult to specify, deformation such as the spatula bending against the
table is not yet maturely implemented in simulators, and contact models such
as point contact have been known to poorly approximate the complex real-world
contact behavior [masterjohn2021discrete] in these settings. In this case,
real environments are out-of-domain (OOD) from simulation, and performing Sys-
ID of simulation parameters might fail to train a useful policy for the real
world due to this inherent irreducible gap.

In this work we take a _task-driven_ approach: instead of trying to align the
simulation with real-world dynamics, we focus on finding simulation parameters
such that the resulting policy optimizes task performance. Such an approach
can lead to policies that achieve high reward in the real world even with an
irreducible sim-to-real gap. We consider settings where the robot has access
to a simulator between iterations of real-world interactions, allowing it to
observe real-world dynamics and update the simulator accordingly with the goal
of improving task performance in reality. Additionally, this approach enhances
the efficiency of real data usage by identifying only task-relevant simulation
parameters and helps trained policies better generalize to OOD (real)
environments.

Statement of Contributions. We propose _AdaptSim_ , an approach for efficient
sim-to-real behavior transfer that interleaves simulation learning with real-
world exploration. AdaptSim iteratively adapts the simulation parameter
distribution used during policy training to maximize performance in a target
environment such as reality. The following is a summary of our contributions:

  * •

We propose an adaptation policy that adjusts a simulation parameter
distribution based on the current task policy’s performance and behavior in
the target environment. Instead of predicting the optimal distribution in one
step, the adaptation incrementally changes the distribution through a small
number of interactions with the target environment. The adaptation policy is
trained using Reinforcement Learning (RL) to maximize the future task reward
achieved.

  * •

We propose a two-phase setup where (i) the adaptation policy is first meta-
trained to adapt to different simulated target environments (randomly sampled
from the entire parameter space in simulation), and (ii) then deployed on the
real environment iteratively. We formalize the meta-learning setting by
introducing the Simulation-Adaptation Contextual Bandit (SA-CB). In the second
phase, the policy uses a small number of interactions with the real
environment in order to infer the optimal change of simulation parameter
distribution for training a task policy to execute in the next iteration. The
same adaptation policy can be re-used for adapting to different real
environments.

  * •

We demonstrate our approach achieving 1-3x asymptotic performance and $\sim$2x
real data efficiency when adapting to different hardware or simulated OOD
environments in three robotic tasks including two that involve contact-rich
manipulation, compared to methods based on Sys-ID and directly training the
task policy in target environments.

Figure 2: AdaptSim consists of two phases: (1) meta-training an adaptation
policy in simulation by maximizing task reward on randomly sampled simulated
target environments; (2) iteratively adapting simulation parameter
distributions based on real trajectories. In simulation, the adaptation policy
$f:(\mathcal{E},\tau)\mapsto\Delta_{\mathcal{E}}$ is trained to adjust the
current simulation parameter distribution based on the task policy’s
trajectory in the target environment; the new distribution is predicted to
achieve better task performance in the next iteration. In the real
environment, the adaptation policy selects a better simulation parameter
distribution using the trained $f$, and the task policy is iteratively fine-
tuned in simulation. The upper-right illustration shows an example of this
process: using only a few real trajectories, the task policy is adapted to
push the bottle closer to the target location (yellow cross).

##  II Related Work

Sim-to-real transfer in robotics has been primarily addressed using Domain
Randomization (DR) techniques [akkaya2019solving, hsu2023sim, tobin2017domain,
sadeghi2016cad2rl, peng2018sim, margolis2022walk] that inject noise in
simulation parameters related to visuals, dynamics, and actuations. A
comprehensive review can be found in [muratore2021robot]. A straightforward
way of designing the simulation parameter setting is to hand-engineer it and
keep it fixed throughout the task policy training, which can be sample-
inefficient and even hinder policy training [mehta2020active]. Below we
summarize techniques that better adapt to real environments.

Sys-ID domain adaptation. Inspired by classical work in Sys-ID
[gautier1988identification, khosla1985parameter], there has been a popular
line of work identifying simulation parameters that match the robot and
environment dynamics in the real environment before task policy training.
BayesSim [ramos2019bayessim] and follow-up work [heiden2021probabilistic,
antonova2022bayesian] applies Bayesian inference to iteratively search for a
posterior distribution of the simulation parameters based on simulation and
real-world trajectories. The inference problem has also been formulated using
RL to minimize trajectory discrepancies [chebotar2019closing]. A different
approach [allevato2020tunenet, ajay2019combining, zeng2020tossingbot] learns a
residual model of dynamics (often parameterized with a neural network) to
match simulation or an ideal physics model with reality. However, all these
methods consider relatively well-modeled environment parameterizations such as
object mass or friction coefficient during planar contact; Sys-ID approaches
have been shown to fail in cases where the simulation does not closely
approximate the real world [muratore2021robot, chi2022iterative]. There is
also work that avoids inferring the full dynamics but adapts with a low-
dimensional latent representation online [kumar2021rma, yu2017preparing,
evans2022context], but the representation is still trained with regression to
match dynamics or simulation parameters. Importantly, the Sys-ID approaches
highlighted above are all task-agnostic; this can lead to poor performance
when trained task policies are sensitive to mismatches in dynamics between
simulation and reality. Chi et al. [chi2022iterative] address the issue by
using simulation to predict changes to trajectories from changes in actions as
an implicit policy, but it requires the environment to be resettable, while
AdaptSim works with randomly initialized object states.

Task-driven domain adaptation. AdaptSim better fits within a different line of
work that aims to find simulation parameters that maximize the task reward in
target environments. Muratore et al. [muratore2021data] apply Bayesian
Optimization (BO) to optimize parameters such as pendulum pole mass and joint
damping coefficient in a real pendulum swing-up task. Other work focus on
adapting to simulated domains only [vuong2019pick, yu2018policy,
ruiz2018learning]. One major drawback of these methods is that they require a
large number of rollouts in target environments (e.g., 700 in
[muratore2021data]), which is very time-consuming for many tasks requiring
human reset. AdaptSim meta-learns adaptation strategies in simulation and
requires only a few real rollouts for inference (_e.g.,_ 20 in our pushing
experiments). Liang et al. [liang2020learning] apply the same task-driven
objective to learn an exploration policy in manipulation tasks, but the task
policy is synthesized using estimated simulation parameters via Sys-ID. Jin et
al. [jin2022task] applies task-drived reduced-order model for dexterous
manipulation tasks, but again the model is identified with Sys-ID and no
vision-based control is involved. Ren et al. [ren2022distributionally] search
for adversarial environments (_e.g.,_ objects) given the current task
performance to robustify the policy, but unlike AdaptSim, the adversarial
metric is measured in simulated domain only without real data.

Learn to search/optimize. Our work involves learning optimization strategies
through meta-learning across a distribution of relevant problems, allowing for
customization to the specific setting and increased sample efficiency
[gomes2021meta, wistuba2018scalable]. Chen et al. [chen2017learning] meta-
learns an RNN optimizer for black-box optimization. Volpp et al.
[volpp2019meta] meta-learns the acquisition function in BO with RL; it is able
to learn new exploration strategies for black-box optimization and tuning
controller gains in sim-to-real transfer. Meta RL trains the task policy
directly to optimize performance in new environments [ren2022leveraging,
johannsmeier2019framework, nagabandi2018neural] — AdaptSim applies meta RL to
optimize simulation parameters instead.

##  III Problem formulation

Environment. In simulation, we consider a space $\Omega$ that parameterizes
the environment with quantities such as friction coefficients and dimensions
of geometric primitives. Let $\mathcal{E}$ denote a distribution of simulation
parameters with support on $\Omega$. We sample an individual environment
$E\in\Omega$ from $\mathcal{E}$ and we denote a real environment $E^{r}$.
Importantly, we do not parameterize the underlying space of $E^{r}$, since it
may be arbitrarily complex or impossible to be fully captured by the
particular $\Omega$ available with existing simulators.

Task Policy. We denote a task policy
$\pi\in\Pi:\mathcal{O}\rightarrow\mathcal{A}$ that maps the robot’s
observation $o_{t}$ to action $a_{t}$. The observation space $\mathcal{O}$ and
the action space $\mathcal{A}$ depend on the task considered.

Trajectory. Running the task policy in a particular environment results in a
state-action trajectory
$\tau(\pi;E):[0,T]\times\Pi\times\mathcal{E}\rightarrow\mathcal{S}\times\mathcal{A}$
with time horizon $T$. The trajectory is also subject to an initial state
distribution. We specify tasks for the robot using a reward function, and let
$R(\tau)\in[0,1]$ denote the normalized cumulative reward accrued by a
trajectory. We let $R(\pi;E)$ denote the reward of running the task policy
$\pi$ in the environment $E$, in expectation over the initial state
distribution.

Goal. Our eventual goal is to find a task policy that maximizes the task
performance in a real environment $E_{r}$:

$\sup_{\pi}\ R(\pi;E^{r}).$ (1)

In this work, we use the distribution $\mathcal{E}$ of simulation parameters
to compactly parameterize the policy $\pi$. This leads to the following bi-
level optimization objective:

$\displaystyle\sup_{\mathcal{E}}\ R(\pi^{*}_{\mathcal{E}};E^{r}),\
\text{where}\ \pi^{*}_{\mathcal{E}}:=$ $\displaystyle\
\sup_{\pi}\underset{E\sim\mathcal{E}}{\mathbb{E}}\ [R(\pi;E)],$ (2)

the optimal task policy for the training distribution $\mathcal{E}$.
Performing the outer level of (2) requires interactions with $E^{r}$ (the real
world); we allow a small budget of such interactions. We emphasize that the
objective above identifies the optimal distribution of simulation parameters
for maximizing task performance, unlike objectives that attempt to match the
dynamics between simulation and reality.

##  IV Approach

###  IV-A Overview

One way to solve (2) is to perform blackbox optimization on $\mathcal{E}$ by
evaluating $R(\pi^{*}_{\mathcal{E}};E^{r})$ [muratore2021data]. However, this
would require a large budget of real trajectories, which is time- and effort-
consuming. AdaptSim instead iteratively _infers_ simulation parameters that
improve task performance by interacting with the real environment a small
number of times. The inference is performed with an “adaptation” policy
$f:(\mathcal{E},\tau)\mapsto\Delta_{\mathcal{E}}$ that adjusts the current
simulation parameter distribution $\mathcal{E}$ by $\Delta_{\mathcal{E}}$
based on $\tau(\pi^{*}_{\mathcal{E}};E^{r})$, the trajectory of the trained
task policy in the real environment.

This approach can be viewed as an iterative adaptation process which predicts
optimal parameter distributions for simulation training based on a small
number of observed real-world trajectories. Using the scooping domain from
Fig. 1 as an example, at each iteration the robot attempts to scoop up pieces
of food in the real world, observes the effects of its attempt, and adjusts
simulation parameters before updating its task policy in simulation.

To eventually achieve adaptation in the real environment, AdaptSim consists of
two phases (Fig. 2):

  1. 1)

Meta-learn the adaptation policy in sim: randomly sample target environments
$E^{s}\in\Omega$ in simulation, and then train $f(\mathcal{E},\tau)$ using RL
to maximize task reward in them, by updating the simulation parameter
distribution (and the corresponding task policies) in iterations.

  2. 2)

Iteratively adapt sim parameters with real data: given a real environment
$E^{r}$ at deployment time, iteratively infer better simulation parameter
distributions for $E^{r}$ using the trained $f$ and a few real trajectories;
the task policy is iteratively fine-tuned in simulation to improve task reward
with the updated parameter distribution.

###  IV-B Phase 1: meta-learning the adaptation policy in sim

In order to correctly infer improved simulation parameters for an unseen real
environment at test time, we first train the adaptation policy to infer better
parameters for many simulated target environments. This phase happens entirely
in simulation. Formally, we model the problem as a partially-observable
contextual bandit [bensoussan1992stochastic]:

###### Definition 1

A Simulation-Adaptation Contextual Bandit (SA-CB) is specified by a tuple of
$(\Omega,\mathcal{T},\mathcal{P},R)$, where

  * •

$\Omega$ is the space of contexts. Each context corresponds to a simulated
target environment $E^{s}$; the context is not directly observable.

  * •

$\mathcal{T}$ is the space of partial observations of the context. Each
observation corresponds to a trajectory observed by running the task policy in
a given context.

  * •

$\mathcal{P}$ is the space of actions. An action corresponds to a choice of
the simulation parameter distribution $\mathcal{E}$.

  * •

$R$ is the reward associated with choosing an action in a particular context
(_i.e.,_ the reward $R(\pi^{*}_{\mathcal{E}};E^{s})$ of the task policy
$\pi^{*}_{\mathcal{E}}$ trained with some simulation parameter distribution
$\mathcal{E}$ when deployed in the target environment).

It may be difficult to infer the optimal $\mathcal{E}\in\mathcal{P}$ using a
single iteration of interactions with the target environment — if the current
task policy fails badly in the target environment, the interaction may reveal
little information. Thus, we iteratively apply incremental changes to
$\mathcal{E}$, with the parameter distribution initialized as
$\mathcal{E}_{i=0}$. Solving the SA-CB (using techniques that we detail
below), we meta-learn an adaptation policy $f(\mathcal{E},\tau)$ to maximize:

$\displaystyle\underset{E^{s}\sim\mathcal{U}_{\Omega}}{\mathbb{E}}\
\underset{\mathcal{E}_{0}\sim\mathcal{U}_{\mathcal{P}}}{\mathbb{E}}\
\sum_{i=0}^{I}\ $ $\displaystyle\gamma^{i}R(\pi^{*}_{\mathcal{E}_{i}};E^{s}),$
(3) where
$\displaystyle\mathcal{E}_{i+1}=\mathcal{E}_{i}+\Delta_{\mathcal{E}_{i}},$
$\displaystyle\Delta_{\mathcal{E}_{i}}=f\big{(}\mathcal{E}_{i},\tau(\pi^{*}_{\mathcal{E}_{i}};E^{s})\big{)},$

and $\mathcal{U}_{\Omega}$ and $\mathcal{U}_{\mathcal{P}}$ are uniform
distributions over $\Omega$ and $\mathcal{P}$ respectively, and $\gamma<1$ is
the discount factor. This is the expected discounted sum of task rewards over
multiple interactions from $i=0$ to the adaptation horizon $I$, over random
sampling of simulated target environment and initial simulation parameter
distribution.

In practice, we are only concerned with the reward if it reaches some minimum
threshold — a bad task policy is not useful. Thus we use a sparse-reward
version of Eq. (3),

$\underset{E^{s}\sim\mathcal{U}_{\Omega}}{\mathbb{E}}\
\underset{\mathcal{E}_{0}\sim\mathcal{U}_{\mathcal{P}}}{\mathbb{E}}\
\sum_{i=0}^{I}\gamma^{i}\mathds{1}\big{(}R(\pi^{*}_{\mathcal{E}_{i}};E^{s})\geq\,\overline{\\!{R}}\big{)}R(\pi^{*}_{\mathcal{E}_{i}};E^{s}),$
(4)

where $\mathds{1}()$ is the indicator function and $\,\overline{\\!{R}}$ is
the sparse-reward threshold. Using a sparse reward also discourages the
adaptation policy from being myopic and getting trapped at a sub-optimal
solution, especially since we use a relatively small $I$ (_e.g.,_ 5-10) in
order to minimize the amount of real data, and use a small discount factor
$\gamma$ ($=0.9$).

Figure 3: AdaptSim uses task-policy trajectories that better reveal task-
relevant information such as scooping dynamics under fast contact (a-1,2),
while simply pushing around the pieces fails to capture it (b).

Task-policy trajectory as observation. We have chosen the task policy
$\pi^{*}_{\mathcal{E}}$ to generate the trajectory observations used by the
adaptation policy. While Sys-ID approaches use arbitrary policies or ones that
generate the most “informative” trajectories in terms of dynamics
[swevers1997optimal], they might not be the most informative in terms of
revealing optimal parameters for task policy training. Our intuition is that
policies that are trained to maximize the task reward are more appropriate
since they better reveal the task-relevant information of the target
environment. In the scooping task, the robot needs to attempt to scoop up the
pieces so it can learn about the effect of the environment on the task (e.g.,
a piece with a flat bottom is generally harder to scoop). Simply pushing the
pieces around does not exhibit the behavior of the pieces under fast contact
(Fig. 3).

Training the adaptation policy using RL. We apply reinforcement learning (RL)
to train the adaptation policy to maximize Eq. (4). In simulation, we collect
$K$ “adaptation trajectories” $\\{\big{(}\mathcal{E}_{i},\
\Delta_{\mathcal{E}_{i}},\ R(\pi^{*}_{\mathcal{E}_{i}};E^{s}),\
\tau(\pi^{*}_{\mathcal{E}_{i}};E^{s})\big{)}\\}_{i=0}^{I}$, whose steps are
saved in a replay buffer $S_{f}$. Since each step involves training the
corresponding task policy $\pi^{*}_{\mathcal{E}_{i}}$, which can be expensive,
we apply off-policy Double Q-Learning [van2016deep] for sample efficiency.
Using this, the adaptation policy outputs the greedy action of a parameterized
Q function,
$f(\mathcal{E},\tau)=\operatorname*{arg\,max}_{\Delta_{\mathcal{E}}}Q(\mathcal{E},\tau;\Delta_{\mathcal{E}})$.
We use $\epsilon$-greedy exploration (at each step, $\epsilon$ probability of
sampling a random action instead of $f(\mathcal{E},\tau)$), with $\epsilon$
initialized at 1 and gradually annealed to 0.

This constitutes the first phase of AdaptSim. Algorithm 1 details the steps
for collecting adaptation trajectories in the inner loop (Line 4-15) and meta-
learning the adaptation policy. We save all distributions (and their
corresponding task policies, omitted in notations for convenience) in a set
$S_{\mathcal{E}}$, which are used again in the second phase. Training the task
policy for each $\mathcal{E}$ is the most computationally heavy component of
Algorithm 1; in Sec. V-C we explain the heuristics applied to allow re-using
task policies between $\mathcal{E}$ that are close to each other in order to
improve computational efficiency.

Algorithm 1 Meta-learning the adaptation policy in sim

1:$(\Omega,\mathcal{T},\mathcal{P},R)$, SA-CB

2:$S_{f}=\varnothing$, replay buffer

3:$S_{\mathcal{E}}=\varnothing$, set of all simulation parameter distributions
(and their task policies) used

4:Initialize $\epsilon\leftarrow 1$

5:for $k\leftarrow 0$ to $K$ do

6: Sample target $E^{s}\sim\mathcal{U}_{\Omega}$ and
$\mathcal{E}_{i=0}\sim\mathcal{U}_{\mathcal{P}}$

7: for $i\leftarrow 0$ to $I$ do

8: Train task policy $\pi^{*}_{\mathcal{E}_{i}}$ with heuristics in Sec.V-C

9: Collect $\tau(\pi^{*}_{\mathcal{E}_{i}};E^{s})$ and
$R(\pi^{*}_{\mathcal{E}_{i}};E^{s})$

10: if $\text{rand}()<\epsilon$ then

11: Choose a random $\Delta_{\mathcal{E}_{i}}$

12: else

13: $\Delta_{\mathcal{E}_{i}}=f(\mathcal{E}_{i},\tau(\cdot;\cdot))$

14: end if

15: Update
$\mathcal{E}_{i+1}\leftarrow\mathcal{E}_{i}+\Delta_{\mathcal{E}_{i}}$

16: Add
$\big{(}\mathcal{E}_{i},\Delta_{\mathcal{E}_{i}},\tau(\cdot;\cdot),R(\cdot;\cdot)\big{)}$
to $S_{f}$

17: Add $\mathcal{E}_{i}$ (and $\pi^{*}_{\mathcal{E}_{i}}$) to
$S_{\mathcal{E}}$

18: end for

19: Train $f$ using Double Q-Learning and $S_{f}$

20: Anneal $\epsilon$ towards $0$

21:end for

22:return $f$, $S_{\mathcal{E}}$

###  IV-C Phase 2: iteratively adapt sim parameters with real data

After meta-training the adaptation policy to find good task policies for a
diverse set of target environments in simulation, we can apply it for
inference and perform adaptation for the real environment $E^{r}$. Algorithm 2
details the iterative process. We apply the same adaptation process as the
inner loop of Algorithm 1 for $I^{r}$ iterations: train the task policy in
simulation, evaluate it in the real environment, and infer the change of
simulation parameters based on real trajectories. We always apply the greedy
action from $f(\mathcal{E},\tau)$ (_i.e.,_ $\epsilon=0$).

Since we have already sampled a large set $S_{\mathcal{E}}$ of parameter
distributions and trained their task policies in Phase 1, we may re-use them
here. At the beginning of Phase 2, we sample $S^{\prime}_{\mathcal{E}}$, a set
of $N$ distributions saved in $S_{\mathcal{E}}$, as the initial distributions
to be adapted independently. Usually we pick $N=2$ considering the trade-off
between number of real trajectories needed and convergence of task performance
(see Appendix LABEL:app:experiment for analysis).

Algorithm 2 Iteratively adapt sim parameters with real data

1:$E^{r}$, real environment

2:$(\Omega,\mathcal{T},\mathcal{P},R)$, SA-CB

3:$f$, adaptation policy trained in Phase 1

4:$S_{f}$, set of simulation parameter distributions (and corresponding task
policies) from Phase 1

5:Sample $S^{\prime}_{f}$ from $S_{f}$

6:for $i\leftarrow 0$ to $I^{r}$ do

7: for $\mathcal{E}_{i}\in S^{\prime}_{f}$ do

8: Train or fine-tune the task policy $\pi^{*}_{\mathcal{E}_{i}}$ in
simulation

9: Collect $\tau(\pi^{*}_{\mathcal{E}_{i}};E^{r})$ and
$R(\pi^{*}_{\mathcal{E}_{i}};E^{r})$ in real

10: Update
$\mathcal{E}_{i+1}\leftarrow\mathcal{E}_{i}+f\big{(}\mathcal{E}_{i},\tau(\cdot;\cdot)\big{)}$

11: end for

12:end for

13:return $\pi^{*}_{\mathcal{E}_{i}}$ with the highest $R(\cdot;E^{r})$

##  V Implementation design choices

Here we elaborate a few design choices, practical considerations, and
implementation details of AdaptSim. Additional details are provided in
Appendix LABEL:app:adaptation-policy.

###  V-A Trajectory observation

First, we remove the action sequence from the task-policy trajectory and keep
the state sequence only. Since the dynamics in real environments can be OOD,
in order to achieve similar high-reward states as in simulated environments,
the robot would need to use some actions not seen during training (or not seen
for the particular state), hindering the adaptation policy to generalize if
action sequence were included in the task policy trajectory. We assume that
the task-relevant _state_ sequence is covered by $\mathcal{T}$ if the task
policy performs reasonably well in the real environment. This choice is also
present in the state-only inverse RL literature [fu2017learning] that
addresses train-test dynamics mismatch. See Fig. LABEL:fig:push-traj-sim-real
and related discussions in Sec. LABEL:subsec:discussions.

We obtain the ground-truth trajectory information in simulation, and estimate
it using either point cloud or color contours with a RGB-D camera in the real
environment.

###  V-B Simulation parameter distribution

In this work, we choose the space $\mathcal{P}$ of possible simulation
parameter distributions in the SA-CB to be that of Gaussian distributions with
mean bounded within $\Omega$ (the space of possible simulation parameter
settings). The diagonal covariance is fixed and small (typically $1\%$ of the
range of the parameter). We choose to adapt only the mean of distributions to
limit the search space and thus save the amount of simulation needed for the
adaptation policy training.

We also use a fixed step size $\delta$ for adapting each simulation parameter,
ranging from $10\%$ to $15\%$ of the parameter range depending on the
dimension of $\Omega$. The set of possible $\Delta_{\mathcal{E}}$ along each
dimension is $\\{\delta,-\delta,0\\}$. Smaller step size is used for lower
dimensional $\Omega$ for more fine-grained adaptation and task policy training
(see Appendix LABEL:app:experiment for sensitivity analysis).

The adaptation policy is parameterized using a Branching Dueling Q-Network
[tavakoli2018action], which outputs the state-action value of choosing any of
the $\\{\delta,-\delta,0\\}$ along each action dimension. It takes in (1) the
vector of the mean of current simulation parameter distribution and (2)
trajectory observation.

Figure 4: Setup of the dynamic pushing and dynamic scooping tasks in both
simulation and reality.

###  V-C Task policy reuse across parameter distributions

Algorithm 1 requires training the task policy for each $\mathcal{E}$, which
can be expensive with the two manipulation tasks. Our intuition is that we can
share the task policy between parameter distributions of close distance, with
the following heuristics:

  * •

Record the total budget (_i.e.,_ number of trajectories), and $j$, the number
of simulation parameter distributions that a task policy has been trained
with.

  * •

Define distance between two parameter distribution $D(\cdot,\cdot)$ such as L2
distance between the mean. If $\mathcal{E}_{i}$ is within a threshold
$\,\overline{\\!{D}}$ from a previously seen distribution, re-use the task
policy. If the policy is already trained with $M_{\text{max}}$ budget total,
do not train again; otherwise train with $\max(M_{\text{min}},\
\alpha^{j-1}M)$ budget, where $\alpha<1$ and $M$ is the budget for training
the policy for the first time.

  * •

If the nearby parameter distribution re-uses a task policy, do not re-use the
same policy again. This prevents the same task policy being used for too many
$\mathcal{E}$.