summaries: This paper proposed to improve existing meta learning algorithms in the presence of task imbalance, class imbalance, and out-of-distribution tasks. Starting from the model-agnostic meta-learning (MAML) algorithm (Finn et al. 2017), to tackle task imbalance, where the number of training examples of varies across different tasks, a task-dependent learning rate decaying factor was learned to be large for large tasks and small for small tasks. In this way, the small task can benefit more from the meta-knowledge and the large task can benefit more from task-specific training. To tackle class imbalance, a class-specific scaling factor was applied to the class-specific gradient. The scaling factor was large for small class and small for large class so that different classes can be treated equally. To tackle the out-of-distribution tasks, a task-dependent variables was learned to emphasize meta-knowledge for the test task similar to training tasks. Additional model parameters are learned through variational inference. Experimental results on benchmark datasets demonstrate the proposed approach outperformed its competing alternatives. Analysis of each component confirm they work as expected.

summaries: The paper proposes a Bayesian approach for meta learning in settings were the tasks might be OOD or have imbalanced class distribution. The proposed approach has 3 task-specific balancing variables with a prior and an inference network. Using an amortized inference scheme, the model unifies the meta-learning objective loss with the lower bound of probabilistic model marginal likelihood.

summaries: This paper introduces a mechanism for gradient-based meta-learning models for few-shot classification to be able to adapt to diverse tasks that are imbalanced and heterogeneous. In particular, each encountered task may have varying numbers of shots (task imbalance) and even within each task, different classes may have different numbers of shots (class imbalance). Further, test tasks might come from a different distribution than the training tasks. They propose to handle this scenario by introducing three new types of variables which control different facets of the degree and type of task adaptation, allowing to decide how much to reuse meta-learned knowledge versus new knowledge acquired from the training set of the given task.

abstractText: While tasks could come with varying the number of instances and classes in realistic settings, the existing meta-learning approaches for few-shot classification assume that the number of instances per task and class is fixed. Due to such restriction, they learn to equally utilize the meta-knowledge across all the tasks, even when the number of instances per task and class largely varies. Moreover, they do not consider distributional difference in unseen tasks, on which the meta-knowledge may have less usefulness depending on the task relatedness. To overcome these limitations, we propose a novel meta-learning model that adaptively balances the effect of the meta-learning and task-specific learning within each task. Through the learning of the balancing variables, we can decide whether to obtain a solution by relying on the meta-knowledge or task-specific learning. We formulate this objective into a Bayesian inference framework and tackle it using variational inference. We validate our Bayesian Task-Adaptive Meta-Learning (Bayesian TAML) on multiple realistic taskand class-imbalanced datasets, on which it significantly outperforms existing meta-learning approaches. Further ablation study confirms the effectiveness of each balancing component and the Bayesian learning framework.

authors:
Luca Bertinetto
Joao F Henriques
Philip HS Torr
Andrea Vedaldi

title: Meta-learning with differentiable closed-form solvers

venue: In ICLR,

authors:
Harrison Edwards
Amos Storkey

title: Towards a neural statistician

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

venue: In ICML,

authors:
Chelsea Finn
Kelvin Xu
Sergey Levine

title: Probabilistic model-agnostic meta-learning

venue: In NeurIPS,

authors:
Sebastian Flennerhag
Andrei A. Rusu
Razvan Pascanu
Francesco Visin
Hujun Yin
Raia Hadsell

title: Meta-learning with warped gradient descent

authors:
Jonathan Gordon
John Bronskill
Matthias Bauer
Sebastian Nowozin
Richard E Turner

title: Metalearning probabilistic inference for prediction

authors:
Sebastian Houben
Johannes Stallkamp
Jan Salmen
Marc Schlipsing
Christian Igel

title: Detection of traffic signs in real-world images: The german traffic sign detection benchmark

venue: In The 2013 international joint conference on neural networks (IJCNN),

authors:
Henry Rowley
Nick Fox-Gieg

title: The quick, draw! – a.i

venue: URL http://quickdraw.withgoogle.com

authors:
Taesup Kim
Jaesik Yoon
Ousmane Dia
Sungwoong Kim
Yoshua Bengio
Sungjin Ahn

title: Bayesian model-agnostic meta-learning

authors:
Diederik P. Kingma
Max Welling

title: Auto encoding variational bayes

authors:
Yoonho Lee
Seungjin Choi

title: Gradient-based meta-learning with learned layerwise metric and subspace

authors:
Zhenguo Li
Fengwei Zhou
Fei Chen
Hang Li

title: Meta-sgd: Learning to learn quickly for few shot learning

venue: arXiv preprint arXiv:1707.09835,

authors:
Subhransu Maji
Esa Rahtu
Juho Kannala
Matthew Blaschko
Andrea Vedaldi

title: Fine-grained visual classification of aircraft

venue: arXiv preprint arXiv:1306.5151,

authors:
Yuval Netzer
Tao Wang
Adam Coates
Alessandro Bissacco
Bo Wu
Andrew Y Ng

title: Reading digits in natural images with unsupervised feature learning

authors:
Maria-Elena Nilsback
Andrew Zisserman

title: Automated flower classification over a large number of classes

venue: Sixth Indian Conference on Computer Vision, Graphics & Image Processing,

authors:
Boris Oreshkin
Pau Rodríguez López
Alexandre Lacoste

title: Tadam: Task dependent adaptive metric for improved few-shot learning

authors:
Siyuan Qiao
Chenxi Liu
Wei Shen
Alan L Yuille

title: Few-shot image recognition by predicting parameters from activations

authors:
Sachin Ravi
Alex Beatson

title: Amortized bayesian meta-learning

authors:
Sachin Ravi
Hugo Larochelle

title: Optimization as a model for few-shot learning

authors:
James Requeima
Jonathan Gordon
John Bronskill
Sebastian Nowozin
Richard E Turner

title: Fast and flexible multi-task classification using conditional neural adaptive processes

venue: NeurIPS,

authors:
Olga Russakovsky
Jia Deng
Hao Su
Jonathan Krause
Sanjeev Satheesh
Sean Ma
Zhiheng Huang
Andrej Karpathy
Aditya Khosla
Michael Bernstein

title: Imagenet large scale visual recognition challenge

venue: International journal of computer vision,

authors:
Andrei A Rusu
Dushyant Rao
Jakub Sygnowski
Oriol Vinyals
Razvan Pascanu
Simon Osindero
Raia Hadsell

title: Meta-learning with latent embedding optimization

authors:
Adam Santoro
Sergey Bartunov
Matthew Botvinick
Daan Wierstra
Timothy Lillicrap

title: Metalearning with memory-augmented neural networks

title: Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-.

venue: hook. PhD thesis, Technische Universität München,

authors:
Jake Snell
Kevin Swersky
Richard Zemel

title: Prototypical networks for few-shot learning

venue: In NIPS,

authors:
Sebastian Thrun
Lorien Pratt (eds

title: Learning to Learn

authors:
Eleni Triantafillou
Tyler Zhu
Vincent Dumoulin
Pascal Lamblin
Kelvin Xu
Ross Goroshin
Carles Gelada
Kevin Swersky
Pierre-Antoine Manzagol
Hugo Larochelle

title: Meta-dataset: A dataset of datasets for learning to learn from few examples

authors:
Oriol Vinyals
Charles Blundell
Timothy Lillicrap
Daan Wierstra

title: Matching networks for one shot learning

authors:
Han Xiao
Kashif Rasul
Roland Vollgraf

title: Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms

venue: arXiv preprint arXiv:1708.07747,

heading: 1 INTRODUCTION

text:
Despite the success of deep learning in many real-world tasks such as visual recognition and machine translation, such good performances are achievable at the availability of large training data, and many fail to generalize well in small data regimes. To overcome this limitation of conventional deep learning, recently, researchers have explored meta-learning (Schmidhuber, 1987; Thrun & Pratt, 1998) approaches, whose goal is to learn a model that generalizes well over distribution of tasks, rather than instances from a single task, in order to utilize the obtained meta-knowledge across tasks to compensate for the lack of training data for each task.
However, so far, most existing meta-learning approaches (Santoro et al., 2016; Vinyals et al., 2016; Snell et al., 2017; Ravi & Larochelle, 2017; Finn et al., 2017; Li et al., 2017) have only targeted an artificial scenario where all tasks participating in the multi-class classification problem have equal number of training instances per class. Yet, this is a highly restrictive setting, as in real-world scenarios, tasks that arrive at the model may have different training instances (task imbalance), and within each task, the number of training instances per class may largely vary (class imbalance). Moreover, the new task may come from a distribution that is different from the task distribution the model has been trained on (out-of-distribution task) (See (a) of Figure 1).
Under such a realistic setting, the meta-knowledge may have a varying degree of utility to each task. Tasks with small number of training data, or close to the tasks trained in meta-training step may want to rely mostly on meta-knowledge obtained over other tasks, whereas tasks that are out-of-distribution or come with more number of training data may obtain better solutions when trained in a task-specific
∗Equal contribution
manner. Furthermore, for multi-class classification, we may want to treat the learning for each class differently to handle class imbalance. Thus, to optimally leverage meta-learning under various imbalances, it would be beneficial for the model to task- and class-adaptively decide how much to use from the meta-learner, and how much to learn specifically for each task and class.
To this end, we propose a novel Bayesian meta-learning framework, which we refer to as Bayesian Task-Adaptive Meta-Learning (Bayesian TAML), that learns variables to adaptively balance the effect of meta- and task-specific learning. Specifically, we first obtain set-representations for each task, which are learned to convey useful statistics about the task or class distribution, such as mean, variance, and cardinality (the number of elements in the set), and then learn the distribution of three balancing variables a function of the set: 1) task-dependent learning rate multiplier, which decides how far away to deviate from the meta-knowledge, when performing task-specific learning. Tasks with higher shots could benefit from taking gradient steps afar, while tasks with few shots may need to stay close to the initial parameter. 2) class-dependent learning rate, which decides how much information to use from each class, to automatically handle class imbalance where the number of instances per class can largely vary. 3) task-dependent modulator for initial model parameter, which modifies the shared initialization for each task, such that each task can decide how much and what to use from the shared initial model parameter and what to ignore based on its set representation. This is especially useful when handling out-of-distribution task, which may need to ignore some of the meta-knowledge.
We validate our model on CIFAR-FS and miniImageNet dataset, as well as a new dataset that consists of heterogeneous datasets, under a scenario where every class in each episode can have any number of shots, that leads to task and class imbalance, and where the dataset at meta-test time is different from that of meta-training time. The experimental results show that our Bayesian TAML significantly improves the performance over the existing approaches under these realistic scenarios. Further analysis of each component reveals that the improvement is due to the effectiveness of the balancing terms for handling task and class imbalance, and out-of-distribution tasks.
To summarize, our contribution in this work is threefold:
• We consider a novel problem of meta-learning under a realistic task distribution, where the number of instances across classes and tasks could largely vary, or the unseen task at the meta-test time is largely different from the seen tasks.
• For effective meta-learning with such imbalances, we propose a Bayesian task-adaptive meta-learning (Bayesian TAML) framework that can adaptively adjust the effect of the meta-learner and the task-specific learner, differently for each task and class.
• We validate our model on realistic imbalanced few-shot classification tasks with a varying number of shots per task and class and show that it significantly outperforms existing meta-learning models.

heading: 2 RELATED WORK

text:
Meta-learning Meta-learning (Schmidhuber, 1987; Thrun & Pratt, 1998) is an approach to learn a model to generalize over a distribution of task. The approaches in general can be categorized into either memory-based, metric-based, and optimization-based methods. A memory-based approach (Santoro et al., 2016) learns to store correct instance and label into the same memory slot and retrieve it later, in a task-generic manner. Metric-based approaches learn a shared metric
space (Vinyals et al., 2016; Snell et al., 2017). Snell et al. (2017) defines the distance between the instance and the class prototype, such that the instances are closer to their correct prototypes than to others. As for optimization-based meta-learning, MAML (Finn et al., 2017) learns a shared initialization parameter that is optimal for any tasks within few gradient steps from the initial parameter. Meta-SGD (Li et al., 2017) improves upon MAML by learning the learning rate differently for each parameter. For effective learning of a meta-learner, meta-learning approaches adopt the episodic training strategy (Vinyals et al., 2016) which trains and evaluates a model over a large number of tasks, which are called meta-training and meta-test phase, respectively. However, existing approaches only consider an artificial scenario which samples the classification of classes with exactly the same number of training instances, both within each episode and across episodes. On the other hand, we consider a more challenging scenario where the number of shots per class and task could vary at each episode, and that the task given at meta-test time could be an out-of-distribution task.
Task-adaptive meta-learning The goal of learning a single meta-learner that works well for all tasks may be overly ambitious and leads to suboptimal performances for each task. Thus recent approaches adopt task-adaptively modified meta-learning models. Oreshkin et al. (2018) proposed to learn the temperature scaling parameter to work with the optimal similarity metric. Qiao et al. (2018) also suggested a model that generates task-specific parameters for the network layers, but it only trains with many-shot classes, and implicitly expects generalization to few-shot cases. Rusu et al. (2019) proposed a network type task-specific parameter producer, and Lee & Choi (2018) proposed to differentiate the network weights into task-shared and task-specific weights. Our model also aims to obtain task-specific parameter for each task, but is rather focused on learning how to balance between the meta-learning and task-/class-specific learning.
Probabilistic meta-learning Recently, a probabilistic version of MAML has been proposed (Finn et al., 2018), where they interpret a task-specific gradient update as a posterior inference process under variational inference framework. Kim et al. (2018) proposed Bayesian MAML with a similar motivation but with a stein variational inference framework and chaser loss. Gordon et al. (2019) proposed a probabilistic meta-learning framework where the parameter for a novel task is rapidly estimated under decision theoretic framework, given a set representation of a task. The motivation behind these works is to represent the inherent uncertainty in few-shot classification tasks. Our model also uses Bayesian modeling, but it focuses on leveraging the uncertainties of the meta-learner and the gradient-direction in order to balance between meta- and task- or class-specific learning.

heading: 3 LEARNING TO BALANCE

text:
We first introduce notations and briefly recap the model-agnostic meta-learning (MAML) by Finn et al. (2017). Suppose a task distribution p(τ) that randomly generates task τ consisting of a training set Dτ = {Xτ ,Yτ} and a test set D̃τ = {X̃τ , Ỹτ}. Then, the goal of MAML is to meta-learn the initial model parameter θ as a meta-knowledge to generalize over the task distribution p(τ), such that we can easily obtain the task-specific predictor θτ in a single (or a few) gradient step from the initial θ. Toward this goal, MAML optimizes the following gradient-based meta-learning objective:
min θ ∑ τ∼p(τ) L(θ − α∇θL(θ;Dτ ); D̃τ ) (1)
where α denotes stepsize and L denotes empirical loss such as negative log-likelihood of observations. Note that by meta-learning the initial point θ, the task-specific predictor θτ = θ − α∇θL(θ;Dτ ) can minimize the test loss L(·; D̃τ ) even with Dτ consisting of only a few samples. We can easily extend the Eq. (1), such that we obtain θτ with more than one inner-gradient steps from the initial θ.
However, the existing MAML framework has the following limitations that prevent the model from efficiently solving real-world problems involving class/task imbalance and out-of-distribution tasks.
1. Class imbalance. MAML does not provide any framework to handle class imbalance within each task. Therefore, classes with large number of training instances (head classes) may dominate the task-specific learning during the inner-gradient steps, yielding low performance on classes with fewer shots (tail classes).
2. Task imbalance. The model has a fixed number of inner-gradient steps and stepsize α across all tasks, which prevents the model from adaptively deciding how much to resort to
the meta-knowledge or how much to learn from the given dataset, depending on the number of the training examples per task.
3. Out-of-distribution tasks. The model assumes that the initial model parameter θ will be equally useful for the unseen tasks, but for unseen tasks that are significantly different from the previously seen tasks, the initial parameter may be less useful.

heading: 3.1 TASK-ADAPTIVE META-LEARNING (TAML)

text:
As shown in Figure 1 for the concepts, we introduce three balancing variables ωτ ,γτ , zτ to tackle each problem mentioned above. How to generate these variables will be described in Section 4. Also, see the experimental section for how to generate the realistic tasks with class and task imbalance.
Tackling class imbalance. To handle class imbalance, we vary the learning rate of class-specific gradient for each inner-optimization step. Specifically, for class c = 1, . . . , C, we introduce a set of class-specific scalars ωτ = (ωτ1 , . . . , ω τ C) ∈ [0, 1]C , which are multiplied to each of the class-specific gradients∇θL(θ;Dτ1 ), . . . ,∇θL(θ;DτC), where Dτc is the set of instances and labels for class c. We expect ωτc to be large for tail-classes with small number of training instances, such that they could be considered more in the inner-optimization steps. We generate ωτ with Softmax function, and denote the input to the function as ω̃τ .
Tackling task imbalance. To control whether the model parameter for the current task stays close to the initial parameter or deviate far from it, we introduce task-dependent learning-rate multipliers γτ = (γτ1 , . . . , γ τ L) ∈ [0,∞)L, such that for each layer l = 1, . . . , L, the learning rate becomes γτ1α, γ τ 2α, . . . , γ τ Lα
1. We expect γτ to be large for large tasks, such that they rely more on taskspecific updates, while small tasks use small γτ to benefit more from the meta-knowledge. To amplify the step-size difference between large and small tasks, we generate γτ with an exponential function, and denote the input to the function as γ̃τ .
Tackling out-of-distribution tasks. Finally, we introduce zτ which modulates the initial parameter θ for each task. We expect zτ to learn to relocate the initial θ to a new starting point, such that out-of-distribution (OOD) tasks can deviate much from the shared initialization θ if the current initialization is suboptimal for the given task. Specifically, we use zτ = 1+ z̃τ for the channel of the convolutional network weights and zτ = z̃τ for the biases, which modify the initial parameter as follows: θ0 ← θ ◦ zτ for the weights and θ0 ← θ + zτ for the biases, where θ0 denotes the new initialization modulated by zτ . We denote this operation as θ ∗ zτ , which is similar to task-dependent modulation of batch normalization parameters (Oreshkin et al., 2018; Requeima et al., 2019).
A unified framework. Finally, we assemble all these components together into a single unified framework. With a slight abuse of notations, the update rule can be written as follows:
θ0 = θ ∗ zτ , (2)
θk = θk−1 − γτ ◦α ◦ C∑ c=1 ωτc∇θk−1L(θk−1;Dτc ) for k = 1, . . . ,K (3)
where α is a multi-dimensional global learning rate vector that is learned (Li et al., 2017), and the multiplication operator ◦ is appropriately defined. The last step θK corresponds to the task-specific predictor θτ .

heading: 3.2 BAYESIAN TASK-ADAPTIVE META-LEARNING

text:
We now introduce the variational inference framework for the input of the three balancing variables, that we previously denote as ω̃τ , γ̃τ , z̃τ . Bayesian modeling allows to incorporate randomness in the posterior of those variables. In MAML framework, such randomness generates the ensemble of diverse task-specific predictors, which allows to effectively exploit the information in the given dataset Dτ . Bayesian modeling also allows to find more robust latent structures of those balancing variables (See Figure 7).
1We initially exponentially decayed the learning rate γ ∈ [0, 1], but found it suboptimal as it prevents the learning rate from growing for many-shot tasks.
Firstly, define Xτ = {xτn} Nτ n=1 and Y τ = {yτn} Nτ n=1 for training, and X̃ τ = {x̃τm} Mτ m=1 and Ỹ τ = {ỹτm} Mτ m=1 for test. Let φ
τ denote the collection of ω̃τ , γ̃τ and z̃τ for uncluttered notation. Then, inspired by Gordon et al. (2019) and Finn et al. (2018), we define the generative process for metalearning framework as follows for each task τ (See Figure 2):
p(Yτ , Ỹτ ,φτ |Xτ , X̃τ ;θ) = p(φτ ) Nτ∏ n=1 p(yτn|xτn,φτ ;θ) Mτ∏ m=1 p(ỹτm|x̃τm,φτ ;θ) (4)
for the complete data likelihood. Note that the deterministic θ is shared across all the tasks.
4 VARIATIONAL INFERENCE
The goal of learning for each task τ is to maximize the conditional loglikelihood of the joint dataset D̃τ and Dτ : log p(Ỹτ ,Yτ |X̃τ ,Xτ ;θ). However, solving it involves the true posterior p(φτ |Dτ , D̃τ ), which is intractable. Thus, we resort to amortized variational inference with a tractable form of approximate posterior q(φτ |Dτ , D̃τ ;ψ) parameterized by ψ. We let the three variables share the same inference network pipeline to minimize the computational cost. Further, similarly to Ravi & Beatson (2019), we drop the dependency on the test dataset D̃τ for the approximate posterior, in order to make the two different pipelines consistent; one for meta-training where we observe the whole test dataset, and the other for meta-testing where the test labels are unknown. The form of our approximate posterior is now q(φτ |Dτ ;ψ). It greatly simplifies the inference framework, while ensuring to have a
valid lower bound of the log evidence. Also, considering that performing the inner-gradient steps with the training datasetDτ automatically maximizes the training log-likelihood in MAML framework, we slightly modify the objective so that the expected log-likelihood term only involves the test examples with the appropriate scaling factor. The resultant form of the approximated lower bound that suits for our meta-learning purpose is as follows:
Lτθ,ψ = Nτ +Mτ Mτ Mτ∑ m=1 Eq(φτ |Dτ ;ψ) [ log p(ỹτm|x̃τm,φτ ;θ) ] −KL[q(φτ |Dτ ;ψ)‖p(φτ )]. (5)
We assume q(φτ |Dτ ;ψ) fully factorizes for each variable and also for each dimension as well: q(φτ |Dτ ;ψ) = ∏ c q(ω̃τc |Dτ ;ψ) ∏ l q(γ̃τl |Dτ ;ψ) ∏ i q(z̃τi |Dτ ;ψ) (6)
where we assume that each single dimension of q(φτ |Dτ ;ψ) follows univariate Gaussian having trainable mean and variance. We also let each dimension of prior p(φτ ) factorize into N (0, 1), such that the KL-divergence can have the especially simple closed form (Kingma & Welling, 2014).
The final form of the meta-training minimization objective with Monte-Carlo (MC) approximation for the expectation in Eq. (5) is as follows:
min θ,ψ
1
Mτ Mτ∑ m=1 1 S S∑ s=1 − log p(ỹτm|x̃τm,φτs ;θ) + 1 Nτ +Mτ KL[q(φτ |Dτ ;ψ)‖p(φτ )]. (7)
where φτs ∼ q(φτ |Dτ ;ψ). We implicitly assume the reparameterization trick for φτ to obtain stable and unbiased gradient estimate w.r.t. ψ (Kingma & Welling, 2014). We set the MC sample size to S = 1 for meta-training for computational efficiency. When meta-testing, we perform MC approximation with sufficiently large sample size (e.g. S = 10):
p(ỹτ∗ |x̃τ∗ ;θ) = Eq[p(ỹτ∗ |x̃τ∗ ,φτ ;θ)] ≈ 1
S S∑ s=1 p(ỹτ∗ |x̃τ∗ ,φτs ;θ), φτs ∼ q(φτ |Dτ ;ψ). (8)
or we may naively approximate by taking the expectation inside for computational efficiency:
Eq[p(ỹτ∗ |x̃τ∗ ,φτ ;θ)] ≈ p(ỹτ∗ |x̃τ∗ ,Eq[φτ ];θ). (9) In the experimental section, we show that Eq. (8) largely outperforms Eq. (9).

heading: 4.1 DATASET ENCODING

text:
The main challenge in modeling our variational distribution q(φτ |Dτ ;ψ) is to decide how to refine the training dataset Dτ into informative representation, which is not trivial. This inference network should capture all the necessary statistical information to recognize any imbalances in the dataset Dτ . Mean-pooling (Edwards & Storkey, 2017) or sum-pooling (Zaheer et al., 2017) is frequently used as a practical set-encoder, where each instance in the set is transformed by the shared nonlinearity, and then averaged or summed together to generate a single vector summarizing the set, followed by an additional nonlinearity. However, for the classification dataset Dτ which is the set of (class) sets, those non-hierarchical pooling methods will perform poorly as they ignore the label information. Therefore, we use a two-layer hierarchical set encoder which first encodes each class as a set of samples and then encodes the set as the set of classes (see the Appendix B for the justification).
However, there is an additional limitation of mean pooling when using it to describe a task: it does not recognize the number of elements in the set2. This could be a critical limitation in recognizing the imbalances in the given task. Therefore we explicitly input the number of elements into the set encoder. Yet, the set cardinality alone is insufficient in capturing the distribution of the dataset. Suppose that we have a set containing replications of a single instance. Then, although the set has only limited information, the set encoding cannot recognize it and will overestimate the information. To prevent this, we encode the variance of the set as well.
Based on this intuition, we define the set encoder StatisticsPooling(·) that generates the concatenation of the set statistics such as mean, variance, and cardinality. We use this encoder to first encode each class, and then the whole set of classes as follows:
vτ = StatisticsPooling ( {NN2 (sc)}Cc=1 ) , sc = StatisticsPooling ( {NN1(x)}x∈Xτc ) for classes c = 1, . . . , C. Xτc is the collection of class c examples in task τ . NN1 and NN2 are some neural networks parameterized by ψ. The vector vτ finally summarizes the entire dataset Dτ for few-shot classification. Note that the class-specific balancing variables ω̃τ1 , . . . , ω̃ τ C are generated from s1, . . . , sC , and other task-specific balancing variables γ̃τ and z̃τ are generated from vτ with a few additional layers (See Figure 3).

heading: 5 EXPERIMENTS

text:
We next validate our method on multiple benchmark datasets with more realistic task distribution.
Datasets We validate our method on the following benchmark datasets. CIFAR-FS: This dataset (Bertinetto et al., 2019) is a variant of CIFAR-100 dataset that consists of 100 general object categories. Each class comes with 600 examples, each of which is a color image that contains 32× 32 pixels. We split the dataset into 64/16/20 classes for training/validation/test. SVHN: This dataset (Netzer et al., 2011) is frequently used as an OOD dataset for CIFAR-10 and CIFAR-100. It consists of 26, 032 color images of 32× 32 pixels, from 10 digits classes. miniImageNet: This dataset (Vinyals et al., 2016) is a subset of the ImageNet dataset (Russakovsky et al., 2015). It consists of total 100 classes, each of which has 600 images resized into 84× 84. We split the dataset into subsets containing 64/16/20 classes for training/validation/test. CUB: This dataset contains 11, 788 images of 200 bird species. We resize the images into 84 × 84. We split the dataset into
2We also found out that sum-pooling is very unstable in encoding set of variable size.
100/50/50 classes for training/validation/test. Since this dataset is fine-grained, we regard it as an out-of-distribution dataset of the coarse-grained miniImageNet dataset.
Realistic task distribution To define realistic task distribution p(τ), we first randomly sample C = 5 classes from the whole set of classes. Then with 0.5 probability, we sample the set size Nc ∼ Unif(1, 50) independently for each class c = 1, . . . , C, in order to simulate class imbalance. On the other hand, with the other 0.5 probability, we again sample the set size Nc ∼ Unif(1, 50) but apply the same single sample to all the classes, in order to simulate task imbalance. We set the number of test (or query) examples to 15 for each class.
Experimental setup We use conventional 4-block convolutional neural networks with 32 channels (Finn et al., 2017). All the baselines and our model become transductive through batch normalization at meta-test time, following Finn et al. (2017). We perform early-stopping for all the baselines and our model with meta-validation set performance. We set the number of inner-gradient steps to 5 for meta-training and 10 for meta-testing, for all the models that take inner-gradient steps. See the Appendix A for more details about the experimental setup.3

heading: 5.1 MAIN RESULTS

text:
Any-shot classification Table 1 shows the results under the realistic task distribution with task and class imbalance. We first observe that Meta-SGD and MT-net outperforms MAML. They precondition the inner-gradients with diagonal or block-diagonal matrix (Flennerhag et al., 2020), which seems to provide extra flexibility to handle task and class imbalance. ABML, one of the recent Bayesian meta-learning models, largely underperforms MAML in this setting, mainly because each taskspecific predictor is excessively regularized by the prior distribution. However, our graphical model in Figure 2 does not directly impose a prior distribution on task-specific predictors, but only on the balancing variables φ. Prototypical Networks perform relatively well on in-distribution (ID) tasks but not on out-of-distribution (OOD) tasks. This is because metric-based models do not take the gradient steps for OOD tasks, such that the novel information in the dataset cannot be effectively incorporated into each task-specific predictor. Proto-MAML has been proposed to take the advantage of both metric-based and gradient-based approach (Triantafillou et al., 2020), but it does not outperform Prototypical Networks in our experiments. Finally, we observe that our Bayesian TAML significantly outperforms all the baselines on both ID and OOD datasets. Bayesian TAML performs especially well on OOD datasets, because when the given task largely differs from the meta-knowledge, the model is able to deviate far from the meta-knowledge; however this is difficult for the other baselines.
Multi-dataset any-shot classification We further test our model under a more challenging setting where tasks could come from a highly heterogeneous dataset (Triantafillou et al., 2020). We meta-
3Code is available at https://github.com/haebeom-lee/l2b
train the models with Aircrafts, Quickdraw and VGG-Flower dataset, and meta-test with the three datasets plus two additional datasets for OOD tasks - Traffic Signs and Fashion-MNIST. We see from Table 2 that Bayesian TAML also largely outperforms baselines in this setting. rototypical Networks perform relatively well for this multi-dataset classification due to their ability of learning a flexible metric space. Proto-MAML brings only marginal improvements on Prototypical Networks since it takes additional gradient steps without considering task dependency. On the other hand, Bayesian TAML effectively combines the advantage of both metric-based and gradient-based approaches by task-dependently modulating the gradient steps to handle task and class imbalance as well. See the Appendix A for the detailed experimental setup.

heading: 5.2 EFFECTIVENESS OF THE BALANCING VARIABLES

text:
We now validate the effectiveness of the three balancing variables to clearly show their effectiveness.
zτ for handling distributional shift. zτ modulates the initial model parameter θ, and decides what and how much to use from the meta-knowledge θ based on the relatedness between the initial θ and the task at hand. We validate the effectiveness of this balancing variable by examining the performance of the baseline methods and our models using the datasets (SVHN, CUB) that are highly heterogeneous from the meta-training datasets (CIFAR-FS, miniImageNet). The results in Table 3 show that our model with only zτ can effectively handle these out-of-distribution (OOD) tasks. We observe in Figure 4 that zτ actually relocates the initial θ far from the initial parameter for these OOD tasks given at meta-test time, with larger displacements for highly heterogeneous tasks (Figure 4, right). This allows the model to either stick to, or deviate from the meta-knowledge based on the similarity between the tasks given at the meta-training and meta-test time. Moreover, the last two rows of Table 3 show that the MC approximation in Eq. (8) largely outperforms the naive approximation in Eq. (9), which suggests that zτ learns very large variance. We conjecture that the role of such random initialization in MAML framework is to increase the effective learning rate for OOD tasks (See Appendix D for the discussion).
γτ and zτ for handling task imbalance. We then examine how the two balancing variables, γτ and zτ handle inter-task imbalance where tasks used for meta-learning have largely different number of instances. Figure 5(a) shows the performance of the baseline models and our models with each of the balancing variables, when the number of instances per task varies from 5 to 2000. We observe that our Bayesian z-TAML and Bayesian γ-TAML largely outperform Meta-SGD, especially by large degree when the number of instances per task is large. Figure 5(b) shows that the effectiveness of zτ largely depends on the number of MC samples used in Eq. (8), demonstrating the importance of incorporating uncertainties in random initializations for handling task imbalance. We further observe from Figure 5(c) that the task-dependent learning rate multiplier γτ rapidly grows as the number of instances per task increases. This agrees with our intuition that larger tasks should take larger inner-gradient steps, to learn more from the given task rather than resorting to meta-knowledge.
ωτ for handling class imbalance. ωτ rescales the class-specific gradients to handle class imbalance where the number of instances per class (i.e. shot) largely varies. Table 4 shows the results under the varying degree of class imbalance across the task distribution. We observe that Bayesian ω-TAML outperforms Meta-SGD, especially by larger degree under higher class imbalance. Notably, Bayesian ω-TAML outperforms a heuristic balancing method which divides each class-specific gradient by the cardinality of each class set (Meta-SGD with 1/N ). The accuracy improvements over Meta-SGD in Figure 6 demonstrate that this heuristic balancing method overly emphasizes the tail classes with few training instances, thereby deteriorating the performance on classes with sufficiently large number of instances. On the other hand, our Bayesian ω-TAML learns the appropriate balancing variables which allow to obtain large performance gains on all classes.
0
0.1
0.2
0.3
0.4
0.5
-2
-1
0
1
2
3
4
5
1~5 6~10 11~20 21~30 31~40 41~50
O m
eg a
A cc
u ra
cy Im
p ro
ve m
en ts
o
ve r
M et
aSG
D (
% )
Number of instances per class
M-SGD with 1/N ω-TAML omega
Figure 6: E[ωτ ] and accuracy improvements over Meta-SGD.

heading: 5.3 MORE ABLATION STUDIES

text:
Effectiveness of Bayesian modeling We further demonstrate the effectiveness of Bayesian modeling by comparing it with the deterministic version of our model (Deterministic TAML), where three balancing variables are no longer stochastic. We apply `2 regularization on the variables with coefficients that are equivalent to the KL-divergence in Eq. (5). The results in Table 5 show that the Bayesian TAML significantly outperforms its deterministic counterpart, especially with larger gap on the OOD task (SVHN). We also see from the last two rows of the table that MC approximation in Eq. (8) is more beneficial for the OOD tasks than for the ID tasks (See Appendix D for the discussion). Figure 7 further shows that the balancing variable γτ for handling task imbalance, more sensitively reacts on Bayesian TAML than on Deterministic TAML, which suggests that Bayesian modeling amplifies the effect of the balancing variables.
Dataset encoding Finally, we perform an ablation study of the proposed task encoding network. The results in Table 6 show that the proposed hierarchical encoding scheme for classification dataset, with set cardinality and variance is significantly more effective than simple mean-pooling methods (Zaheer et al., 2017; Edwards & Storkey, 2017)4.

heading: 6 CONCLUSION

text:
We propose Bayesian TAML that learns to balance the effect of meta-learning and task-adaptive learning, to consider meta-learning under a more realistic task distribution where each task and class can have varying number of instances. Specifically, we encode the dataset for each task into
4We also found out that using set cardinality is more effective than utilizing higher-order statistics such as skewness or kurtosis (See the Appendix C for further ablation study).
hierarchical representations, and use it to modulate the original parameter, learning rate, and the classspecific gradients. We use a Bayesian framework to infer the posterior of these balancing variables, and propose an effective variational inference framework to solve for them. Our model outperforms existing meta-learning methods when validated on imbalanced few-shot classification tasks. Further analysis of each balancing variable shows that each variable effectively handles task imbalance, class imbalance, and out-of-distribution tasks. We believe that our work makes a meaningful step toward application of meta-learning to real-world problems.
Acknowledgements This work was supported by Google AI Focused Research Award, Samsung Research Funding Center of Samsung Electronics (No. SRFC-IT1502-51), the Engineering Research Center Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT (NRF-2018R1A5A1059921), and the Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.20190-00075, and Artificial Intelligence Graduate School Program (KAIST)).

heading: A EXPERIMENTAL SETUP

text:
A.1 BASELINES
We describe the baseline models and our task-adaptive learning to balance model.
1) MAML. The Model-Agnostic Meta-Learning (MAML) model by Finn et al. (2017), which aims to learn the global initial model parameter, from which we can take a few gradient steps to get task-specific predictors.
2) ABML. This model interprets MAML under hierarchical Bayesian framework, but they propose to share and amortize the inference rules across both global initial parameters as well as the task-specific parameters.
3) MT-net. A gradient-based meta-learning model proposed by Lee & Choi (2018). The model obtains a task-specific parameter only w.r.t. a subset of the whole dimension (M-net), followed by a linear transformation to learn a metric space (T-net).
4) Meta-SGD. A base MAML with the learnable learning-rate vector (without any restriction on sign) element-wisely multiplied to each step inner-gradient (Li et al., 2017).
5) Prototypical Networks. A metric-based few-shot classification model proposed by Snell et al. (2017). The model learns a metric space based on Euclidean distance between class prototypes and query embeddings.
6) Proto-MAML. A variant of MAML (Triantafillou et al., 2020) that replaces the initialization of the final fully-connected layer matrix with the equivalent one of the Prototypical Networks (Snell et al., 2017). This model combines the advantage of both metric-based and gradient-based approach. We set α to 0.0005 for any-shot classification and 0.01 for multi-dataset experiments.
7) Bayesian TAML. Our model that can adaptively balance between meta- and task-specific learners for each task and class.
A.2 ANY-SHOT CLASSIFICATION.
We describe more detailed experimental settings for any-shot classification. For MAML, ABML and MT-NET, we set the inner-gradient stepsize α to 0.5 for CIFAR-FS / SVHN, and 0.1 for miniImageNet / CUB, after searching the range α ∈ {0.01, 0.05, 0.1, 0.5}. CIFAR-FS and SVHN: We meta-train all models for total 50K iterations with meta-batch size set to 4. The outer learning rate is set to 0.001 for all the baselines and our models.
miniImageNet and CUB: We meta-train all models for total 80K iterations with meta-batch size set to 1. The outer learning rate is set to 0.0001 for all the baselines and our models.
A.3 MULTI-DATASET CLASSIFICATION
For multi-dataset classification, we construct a subset of the whole collection of the Meta-Dataset (Triantafillou et al., 2020). We resize the images into 32× 32 pixels. For each task, We uniformly select one dataset among Aircrafts, Quickdraw and VGG-Flower and randomly sample 10 classes from the dataset. Then we sample instances from each class with the number of instances per class ranging from 1 to 50. The number of test instances is equally set to 15 for all classes. At meta-test time, we use the three datasets plus two more out-of-distribution datasets - Traffic Signs and Fashion-MNIST. For MAML, ABML and MT-NET, we set the inner-gradient stepsize α to 0.5. We set the number of classes for each task to 10, meta-batch size to 3, meta-training iterations to 60K, and outer learning rate to 0.001 for all models.
Aircraft: We split this dataset (Maji et al., 2013) into 70/15/15 classes for meta- training/validation/test with 100 examples for each class.
Quick Draw: We split this dataset (Jonas Jongejan & Fox-Gieg, 2016) into 241/52/52 classes for meta- training/validation/test with randomly sampled 200 examples for each class.
VGG Flower: This dataset (Nilsback & Zisserman, 2008) contains 40 between 258 images for each class and we split this dataset into 71/16/15 classes for meta- training/validation/test.
Traffic Signs: This dataset (Houben et al., 2013) has only test set consisting of 43 classes. Each class has 900 examples.
Fashion-MNIST: We use the test set of Fashion-MNIST (Xiao et al., 2017) for meta-testing. This dataset has 10 classes with 1000 examples per class.
A.4 INFERENCE NETWORK ARCHITECTURE
We describe the network architecture of the inference network that takes a classification dataset as an input and generates three balancing variables as output.
Shared encoder NN1 : Xτ1 , . . . ,XτC → s1, . . . ,xC 3 × 3 Conv2d with 10-dim and ReLU 2 × 2 Max pool 3 × 3 Conv2d with 10-dim and ReLU 2 × 2 Max pool fc layer with 64-dim Statistics Pooling for each class c = 1, . . . , C. fc layer with 4-dim (across the statistics) and ReLU
Shared encoder NN2 : s1, . . . , sC → vτ
fc layer with 128-dim and ReLU fc layer with 32-dim Statistics Pooling over all the class representations fc layer with 4-dim (across the statistics) and ReLU
Generate ωτ : sτ1 , . . . , sτC → (µτω1 , σ τ ω1), . . . , (µ τ ωC , σ τ ωC )
fc layer with 64-dim and ReLU fc layer to generate (µτωc , σ τ ωc) for each class c = 1, . . . , C.
Generate γτ : vτ → (µτγ1 , σ τ γ1), . . . , (µ τ γL , σ τ γL)
fc layer with 64-dim and ReLU fc layer to generate (µτγl , σ τ γl ) for each layer l = 1, . . . , L.
Generate zτ : vτ → µτz ,στz fc layer with 64-dim and ReLU fc layer to generate (µτz ,σ τ z ) for the output channels

heading: B JUSTIFICATION FOR SET-OF-SETS STRUCTURE.

text:
Based on the previous justification of DeepSets (Zaheer et al., 2017), we can easily justify the Set-ofSets structure proposed in the main paper as well, in terms of the two-level permutation invariance properties required for any classification dataset. The main theorem of DeepSets is:
Theorem 1. A function f operating on a set X ∈ X is a valid set function (i.e. permutation invariant), iff it can be decomposed as f(X) = ρ2( ∑ x∈X ρ1(x)), where ρ1 and ρ2 are appropriate nonfcities.
See (Zaheer et al., 2017) for the proof. Here we apply the same argument twice as follows.
1. A function f operating on a set of representations {s1, . . . , sC} (we assume each sc is an output from a shared function g) is a valid set function (i.e. permutation invariant w.r.t. the order of {s1, . . . , sC}), iff it can be decomposed as f({s1, . . . , sC}) = ρ2( ∑C c=1 ρ1(sc))
with appropriate nonfcities ρ1 and ρ2.
2. A function g operating on a set of examples {xc,1, . . . ,xc,N} is a valid set function (i.e. permutation invariant w.r.t. the order of {xc,1, . . . ,xc,N}) iff it can be decomposed as g({xc,1, . . . ,xc,N}) = ρ4( ∑N i=1 ρ3(xc,i)) with appropriate nonfcities ρ3 and ρ4.
Inserting sc = g({xc,1, . . . ,xc,N}) into the expression of f , we arrive at the following valid composite function operating on a set of sets:
f ( {g ({xc,1, . . . ,xc,N})}Cc=1 ) = ρ2 ( C∑ c=1 ρ1 ( ρ4 ( N∑ i=1 ρ3 (xc,i) ))) (10)
Let F denote the composite of f and (multiple) g and let NN2 denote the composite of ρ1 and ρ4. Further define NN1 := ρ3 and NN3 := ρ2. Then, we have
F ({ {x1,1, . . . ,x1,N}, . . . , {xC,1, . . . ,xC,N} }) = NN3 ( C∑ c=1 NN2 ( N∑ i=1 NN1 (xc,i) )) (11)
where C is the number of classes and N is the number of examples per class. See Section A.4 for the actual encoder structure.

heading: C ABLATION STUDY ON HIGHER-ORDER STATISTICS

text:
While our framework does not place any restriction on selecting the statistics for set encoding, we perform further ablation study on the effectiveness of higher-order statistics for our specific experimental setting. We see from Table 7 that the higherorder statistics such as element-wise sample skewness and kurtosis improve the performance given sample mean and diagonal covariance. However, the set cardinality seems more effective than those
statistics as it could be the most direct and relevant criteria for detecting imbalances in a set.

heading: D ANALYSIS ON RANDOM INITIALIZATION FOR SOLVING

text:
OUT-OF-DISTRIBUTION TASKS
In this section, we analyze the effect of the variance of zτ for solving OOD tasks, that randomizes the MAML initialization parameter θ. Define Eq(zτ |Dτ ;ψ)[θ0] = θ ∗ Eq[zτ ], the initialization parameter modulated by the posterior mean of zτ . Then, we measure the `2 distance from Eq[θ0] to the two different versions of the final task-specific parameter θτ after taking 10 gradient steps, in order to see how the posterior variance of zτ affects the final task-specific parameter θτ as a function of zτ :
• θτ (Eq[zτ ]]): Task-specific predictor θτ obtained without the variance of zτ , such that the expectation is taken before the inner-gradient steps. • Ez[θτ (zτ )]: Task-specific predictor θτ obtained with
the variance in the random initialization, such that the expectation is taken outside of the gradient steps. We evaluate the expectation with MC approximation, having the ensemble of the diverse task-specific predictor samples θτ1 , . . . ,θ τ S (we use S = 50).
Figure 8 suggests that the role of the random initialization is to increase the effective learning rate for the OOD tasks. We see from the left bar graph that if we do not consider variance in the initialization (θτ (Eq[zτ ])), the OOD tasks deviate relatively less than the ID tasks (4.61 vs. 5.18), although it
should deviate much considering the distributional discrepancy. On the other hand, if we incorporate the random initialization to obtain task-specific parameter (Eq[θτ (zτ )]), OOD tasks can deviate further from the initialization (4.61→ 5.11). It directly results in the performance gain because the task-specific learner can exploit more of the information in the OOD tasks.

title: OUT-OF-DISTRIBUTION TASKS


summaries: The authors propose SQUIL, an off-policy imitation learning (IL) algorithm which attempts to overcome the classic drift problems of behavioral cloning (BC). The idea is to reduce IL to a standard RL problem with a reward that incentivizes the agent to take expert actions in states observed in the demonstrations. The algorithm is tested on both image-based tasks (Atari) and continuous control tasks (Mujoco) and shown effective against GAIL and a simple supervised BC approach.

summaries: This paper proposes an imitation learning approach via reinforcement learning. The imitation learning problem is transformed into an RL problem with a reward of +1 for matching an expert's action in a state and a reward of 0 for failing to do so. This encourages the agent to return to "known" states from out-of-distribution states and alleviates the problem of compounding errors. The authors derive an interpretation of their approach as regularized behavior cloning. Furthermore, they empirically evaluate their approach on a set of imitation learning problems, showing strong performance. The authors also stress the easy implementability of their approach within standard RL methods.

summaries: This paper proposes a simple method for imitation learning that is competitive with GAIL.  The approach, Soft Q Imitation Learning (SQIL), utilizes Soft Q Learning, and defines high rewards as faithfully following the demonstrations and low rewards as deviating from them.  Because SQIL is off-policy, it can utilize replay buffers to accelerate sample efficiency.  One can also interpret SQIL as a regularized version of behavioral cloning.

abstractText: Learning to imitate expert behavior from demonstrations can be challenging, especially in environments with high-dimensional, continuous observations and unknown dynamics. Supervised learning methods based on behavioral cloning (BC) suffer from distribution shift: because the agent greedily imitates demonstrated actions, it can drift away from demonstrated states due to error accumulation. Recent methods based on reinforcement learning (RL), such as inverse RL and generative adversarial imitation learning (GAIL), overcome this issue by training an RL agent to match the demonstrations over a long horizon. Since the true reward function for the task is unknown, these methods learn a reward function from the demonstrations, often using complex and brittle approximation techniques that involve adversarial training. We propose a simple alternative that still uses RL, but does not require learning a reward function. The key idea is to provide the agent with an incentive to match the demonstrations over a long horizon, by encouraging it to return to demonstrated states upon encountering new, out-of-distribution states. We accomplish this by giving the agent a constant reward of r = +1 for matching the demonstrated action in a demonstrated state, and a constant reward of r = 0 for all other behavior. Our method, which we call soft Q imitation learning (SQIL), can be implemented with a handful of minor modifications to any standard Q-learning or off-policy actor-critic algorithm. Theoretically, we show that SQIL can be interpreted as a regularized variant of BC that uses a sparsity prior to encourage long-horizon imitation. Empirically, we show that SQIL outperforms BC and achieves competitive results compared to GAIL, on a variety of image-based and low-dimensional tasks in Box2D, Atari, and MuJoCo. This paper is a proof of concept that illustrates how a simple imitation method based on RL with constant rewards can be as effective as more complex methods that use learned rewards.

authors:
Brenna D Argall
Sonia Chernova
Manuela Veloso
Brett Browning

title: A survey of robot learning from demonstration

venue: Robotics and autonomous systems,

authors:
Marc G Bellemare
Yavar Naddaf
Joel Veness
Michael Bowling

title: The arcade learning environment: An evaluation platform for general agents

venue: Journal of Artificial Intelligence Research,

authors:
Giuseppe Calafiore
Fabrizio Dabbene

title: Probabilistic and randomized methods for design under uncertainty

authors:
Chelsea Finn
Paul Christiano
Pieter Abbeel
Sergey Levine

title: A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models

venue: arXiv preprint arXiv:1611.03852,

authors:
Chelsea Finn
Sergey Levine
Pieter Abbeel

title: Guided cost learning: Deep inverse optimal control via policy optimization

venue: In International Conference on Machine Learning,

authors:
Justin Fu
Katie Luo
Sergey Levine

title: Learning robust rewards with adversarial inverse reinforcement learning

venue: arXiv preprint arXiv:1710.11248,

authors:
Yang Gao
Ji Lin
Fisher Yu
Sergey Levine
Trevor Darrell

title: Reinforcement learning from imperfect demonstrations

venue: arXiv preprint arXiv:1802.05313,

authors:
Ian Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron Courville
Yoshua Bengio

title: Generative adversarial nets

venue: In Advances in neural information processing systems,

authors:
David Ha
Jürgen Schmidhuber

title: Recurrent world models facilitate policy evolution

venue: arXiv preprint arXiv:1809.01999,

authors:
Tuomas Haarnoja
Haoran Tang
Pieter Abbeel
Sergey Levine

title: Reinforcement learning with deep energybased policies

venue: arXiv preprint arXiv:1702.08165,

authors:
Tuomas Haarnoja
Aurick Zhou
Pieter Abbeel
Sergey Levine

title: Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor

venue: arXiv preprint arXiv:1801.01290,

authors:
Peter Henderson
Riashat Islam
Philip Bachman
Joelle Pineau
Doina Precup
David Meger

title: Deep reinforcement learning that matters

venue: In Thirty-Second AAAI Conference on Artificial Intelligence,

authors:
Todd Hester
Matej Vecerik
Olivier Pietquin
Marc Lanctot
Tom Schaul
Bilal Piot
Dan Horgan
John Quan
Andrew Sendonaris
Gabriel Dulac-Arnold

title: Deep q-learning from demonstrations

venue: arXiv preprint arXiv:1704.03732,

authors:
Jonathan Ho
Stefano Ermon

title: Generative adversarial imitation learning

venue: In Advances in Neural Information Processing Systems,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Ilya Kostrikov
Kumar Krishna Agrawal
Debidatta Dwibedi
Sergey Levine
Jonathan Tompson

title: Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning

venue: In International Conference on Learning Representations,

authors:
Karol Kurach
Mario Lucic
Xiaohua Zhai
Marcin Michalski
Sylvain Gelly

title: The gan landscape: Losses, architectures, regularization, and normalization

venue: arXiv preprint arXiv:1807.04720,

title: Reinforcement learning and control as probabilistic inference: Tutorial and review

venue: arXiv preprint arXiv:1805.00909,

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Andrei A Rusu
Joel Veness
Marc G Bellemare
Alex Graves
Martin Riedmiller
Andreas K Fidjeland
Georg Ostrovski

title: Human-level control through deep reinforcement learning

authors:
Andrew Y Ng
Stuart J Russell

title: Algorithms for inverse reinforcement learning

venue: In Icml, pp

authors:
Bilal Piot
Matthieu Geist
Olivier Pietquin

title: Learning from demonstrations: Is it worth estimating a reward function

venue: In Joint European Conference on Machine Learning and Knowledge Discovery in Databases,

title: Boosted and reward-regularized classification for apprenticeship learning

venue: In Proceedings of the 2014 international conference on Autonomous agents and multi-agent systems,

title: Efficient training of artificial neural networks for autonomous navigation

venue: Neural Computation,

authors:
Stéphane Ross
Geoffrey Gordon
Drew Bagnell

title: A reduction of imitation learning and structured prediction to no-regret online learning

venue: In Proceedings of the fourteenth international conference on artificial intelligence and statistics,

authors:
Fumihiro Sasaki
Tetsuya Yohira
Atsuo Kawaguchi

title: Sample efficient imitation learning for continuous control

authors:
Emanuel Todorov
Tom Erez
Yuval Tassa

title: Mujoco: A physics engine for model-based control

venue: In Intelligent Robots and Systems (IROS),

authors:
Ruohan Wang
Carlo Ciliberto
Pierluigi Amadori
Yiannis Demiris

title: Random expert distillation: Imitation learning via expert policy support estimation

authors:
Markus Wulfmeier
Peter Ondruska
Ingmar Posner

title: Maximum entropy deep inverse reinforcement learning

venue: arXiv preprint arXiv:1507.04888,

authors:
Brian D Ziebart
J Andrew Bagnell
Anind K Dey

title: Modeling interaction via the principle of maximum causal entropy

venue: In Proceedings of the 27th International Conference on International Conference on Machine Learning,

title: The mechanism for learning terminal rewards

heading: 1 INTRODUCTION

text:
Many sequential decision-making problems can be tackled by imitation learning: an expert demonstrates near-optimal behavior to an agent, and the agent attempts to replicate that behavior in novel situations (Argall et al., 2009). This paper considers the problem of training an agent to imitate an expert, given expert action demonstrations and the ability to interact with the environment. The agent does not observe a reward signal or query the expert, and does not know the state transition dynamics.
Standard approaches based on behavioral cloning (BC) use supervised learning to greedily imitate demonstrated actions, without reasoning about the consequences of actions (Pomerleau, 1991). As a result, compounding errors cause the agent to drift away from the demonstrated states (Ross et al., 2011). The problem with BC is that, when the agent drifts and encounters out-of-distribution states, the agent does not know how to return to the demonstrated states. Recent methods based on inverse reinforcement learning (IRL) overcome this issue by training an RL agent not only to imitate demonstrated actions, but also to visit demonstrated states (Ng et al., 2000; Wulfmeier et al., 2015; Finn et al., 2016b; Fu et al., 2017). This is also the core idea behind generative adversarial imitation learning (GAIL) (Ho & Ermon, 2016), which implements IRL using generative adversarial
networks (Goodfellow et al., 2014; Finn et al., 2016a). Since the true reward function for the task is unknown, these methods construct a reward signal from the demonstrations through adversarial training, making them difficult to implement and use in practice (Kurach et al., 2018).
The main idea in this paper is that the effectiveness of adversarial imitation methods can be achieved by a much simpler approach that does not require adversarial training, or indeed learning a reward function at all. Intuitively, adversarial methods encourage long-horizon imitation by providing the agent with (1) an incentive to imitate the demonstrated actions in demonstrated states, and (2) an incentive to take actions that lead it back to demonstrated states when it encounters new, out-ofdistribution states. One of the reasons why adversarial methods outperform greedy methods, such as BC, is that greedy methods only do (1), while adversarial methods do both (1) and (2). Our approach is intended to do both (1) and (2) without adversarial training, by using constant rewards instead of learned rewards. The key idea is that, instead of using a learned reward function to provide a reward signal to the agent, we can simply give the agent a constant reward of r = +1 for matching the demonstrated action in a demonstrated state, and a constant reward of r = 0 for all other behavior.
We motivate this approach theoretically, by showing that it implements a regularized variant of BC that learns long-horizon imitation by (a) imposing a sparsity prior on the reward function implied by the imitation policy, and (b) incorporating information about the state transition dynamics into the imitation policy. Intuitively, our method accomplishes (a) by training the agent using an extremely sparse reward function – +1 for demonstrations, 0 everywhere else – and accomplishes (b) by training the agent with RL instead of supervised learning.
We instantiate our approach with soft Q-learning (Haarnoja et al., 2017) by initializing the agent’s experience replay buffer with expert demonstrations, setting the rewards to a constant r = +1 in the demonstration experiences, and setting rewards to a constant r = 0 in all of the new experiences the agent collects while interacting with the environment. Since soft Q-learning is an off-policy algorithm, the agent does not necessarily have to visit the demonstrated states in order to experience positive rewards. Instead, the agent replays the demonstrations that were initially added to its buffer. Thus, our method can be applied in environments with stochastic dynamics and continuous states, where the demonstrated states are not necessarily reachable by the agent. We call this method soft Q imitation learning (SQIL).
The main contribution of this paper is SQIL: a simple and general imitation learning algorithm that is effective in MDPs with high-dimensional, continuous observations and unknown dynamics. We run experiments in four image-based environments – Car Racing, Pong, Breakout, and Space Invaders – and three low-dimensional environments – Humanoid, HalfCheetah, and Lunar Lander – to compare SQIL to two prior methods: BC and GAIL. We find that SQIL outperforms BC and achieves competitive results compared to GAIL. Our experiments illustrate two key benefits of SQIL: (1) that it can overcome the state distribution shift problem of BC without adversarial training or learning a reward function, which makes it easier to use, e.g., with images, and (2) that it is simple to implement using existing Q-learning or off-policy actor-critic algorithms.

heading: 2 SOFT Q IMITATION LEARNING

text:
SQIL performs soft Q-learning (Haarnoja et al., 2017) with three small, but important, modifications: (1) it initially fills the agent’s experience replay buffer with demonstrations, where the rewards are set to a constant r = +1; (2) as the agent interacts with the world and accumulates new experiences, it adds them to the replay buffer, and sets the rewards for these new experiences to a constant r = 0; and (3) it balances the number of demonstration experiences and new experiences (50% each) in each sample from the replay buffer.1 These three modifications are motivated theoretically in Section 3, via an equivalence to a regularized variant of BC. Intuitively, these modifications create a simple reward structure that gives the agent an incentive to imitate the expert in demonstrated states, and to take actions that lead it back to demonstrated states when it strays from the demonstrations.
1 SQIL resembles the Deep Q-learning from Demonstrations (DQfD) (Hester et al., 2017) and Normalized Actor-Critic (NAC) algorithms (Gao et al., 2018), in that it initially fills the agent’s experience replay buffer with demonstrations. The key difference is that DQfD and NAC are RL algorithms that assume access to a reward signal, while SQIL is an imitation learning algorithm that does not require an extrinsic reward signal from the environment. Instead, SQIL automatically constructs a reward signal from the demonstrations.
Algorithm 1 Soft Q Imitation Learning (SQIL) 1: Require λsamp ∈ R≥0,Ddemo 2: Initialize Dsamp ← ∅ 3: while Qθ not converged do 4: θ ← θ − η∇θ(δ2(Ddemo, 1) + λsampδ2(Dsamp, 0)) {See Equation 1} 5: Sample transition (s, a, s′) with imitation policy π(a|s) ∝ exp (Qθ(s, a)) 6: Dsamp ← Dsamp ∪ {(s, a, s′)} 7: end while
Crucially, since soft Q-learning is an off-policy algorithm, the agent does not necessarily have to visit the demonstrated states in order to experience positive rewards. Instead, the agent replays the demonstrations that were initially added to its buffer. Thus, SQIL can be used in stochastic environments with high-dimensional, continuous states, where the demonstrated states may never actually be encountered by the agent.
SQIL is summarized in Algorithm 1, where Qθ is the soft Q function, Ddemo are demonstrations, δ2 is the squared soft Bellman error,
δ2(D, r) , 1|D| ∑
(s,a,s′)∈D
( Qθ(s, a)− ( r + γ log (∑ a′∈A exp (Qθ(s ′, a′)) )))2 , (1)
and r ∈ {0, 1} is a constant reward.2 The experiments in Section 4 use a convolutional neural network or multi-layer perceptron to model Qθ, where θ are the weights of the neural network. Section A.3 in the appendix contains additional implementation details, including values for the hyperparameter λsamp; note that the simple default value of λsamp = 1 works well across a variety of environments.
As the imitation policy in line 5 of Algorithm 1 learns to behave more like the expert, a growing number of expert-like transitions get added to the bufferDsamp with an assigned reward of zero. This causes the effective reward for mimicking the expert to decay over time. Balancing the number of demonstration experiences and new experiences (50% each) sampled for the gradient step in line 4 ensures that this effective reward remains at least 1/(1 + λsamp), instead of decaying to zero. In practice, we find that this reward decay does not degrade performance if SQIL is halted once the squared soft Bellman error loss converges to a minimum (e.g., see Figure 8 in the appendix). Note that prior methods also require similar techniques: both GAIL and adversarial IRL (AIRL) (Fu et al., 2017) balance the number of positive and negative examples in the training set of the discriminator, and AIRL tends to require early stopping to avoid overfitting.

heading: 3 INTERPRETING SQIL AS REGULARIZED BEHAVIORAL CLONING

text:
To understand why SQIL works, we sketch a surprising theoretical result: SQIL is equivalent to a variant of behavioral cloning (BC) that uses regularization to overcome state distribution shift.
BC is a simple approach that seeks to imitate the expert’s actions using supervised learning – in particular, greedily maximizing the conditional likelihood of the demonstrated actions given the demonstrated states, without reasoning about the consequences of actions. Thus, when the agent makes small mistakes and enters states that are slightly different from those in the demonstrations, the distribution mismatch between the states in the demonstrations and those actually encountered by the agent leads to compounding errors (Ross et al., 2011). We show that SQIL is equivalent to augmenting BC with a regularization term that incorporates information about the state transition dynamics into the imitation policy, and thus enables long-horizon imitation.

heading: 3.1 PRELIMINARIES

text:
Maximum entropy model of expert behavior. SQIL is built on soft Q-learning, which assumes that expert behavior follows the maximum entropy model (Ziebart et al., 2010; Levine, 2018). In
2Equation 1 assumes discrete actions, but SQIL can also be used with continuous actions, as shown in Section 4.3.
an infinite-horizon Markov Decision Process (MDP) with a continuous state space S and discrete action space A,3 the expert is assumed to follow a policy π that maximizes reward R(s, a). The policy π forms a Boltzmann distribution over actions,
π(a|s) , exp (Q(s, a))∑ a′∈A exp (Q(s, a ′)) , (2)
where Q is the soft Q function. The soft Q values are a function of the rewards and dynamics, given by the soft Bellman equation,
Q(s, a) , R(s, a) + γEs′ [ log (∑ a′∈A exp (Q(s′, a′)) )] . (3)
In our imitation setting, the rewards and dynamics are unknown. The expert generates a fixed set of demonstrations Ddemo, by rolling out their policy π in the environment and generating state transitions (s, a, s′) ∈ Ddemo. Behavioral cloning (BC). Training an imitation policy with standard BC corresponds to fitting a parametric model πθ that minimizes the negative log-likelihood loss,
`BC(θ) , ∑
(s,a)∈Ddemo
− log πθ(a|s). (4)
In our setting, instead of explicitly modeling the policy πθ, we can represent the policy π in terms of a soft Q function Qθ via Equation 2:
π(a|s) , exp (Qθ(s, a))∑ a′∈A exp (Qθ(s, a ′)) . (5)
Using this representation of the policy, we can train Qθ via the maximum-likelihood objective in Equation 4:
`BC(θ) , ∑
(s,a)∈Ddemo
− ( Qθ(s, a)− log (∑ a′∈A exp (Qθ(s, a ′)) )) . (6)
However, optimizing the BC loss in Equation 6 does not in general yield a valid soft Q function Qθ – i.e., a soft Q function that satisfies the soft Bellman equation (Equation 3) with respect to the dynamics and some reward function. The problem is that the BC loss does not incorporate any information about the dynamics into the learning objective, so Qθ learns to greedily assign high values to demonstrated actions, without considering the state transitions that occur as a consequence of actions. As a result, Qθ may output arbitrary values in states that are off-distribution from the demonstrations Ddemo. In Section 3.2, we describe a regularized BC algorithm that adds constraints to ensure that Qθ is a valid soft Q function with respect to some implicitly-represented reward function, and further regularizes the implicit rewards with a sparsity prior. In Section 3.3, we show that this approach recovers an algorithm similar to SQIL.

heading: 3.2 REGULARIZED BEHAVIORAL CLONING

text:
Under the maximum entropy model described in Section 3.1, expert behavior is driven by a reward function, a soft Q function that computes expected future returns, and a policy that takes actions with high soft Q values. In the previous section, we used these assumptions to represent the imitation policy in terms of a model of the soft Q function Qθ (Equation 5). In this section, we represent the reward function implicitly in terms of Qθ, as shown in Equation 7. This allows us to derive SQIL as a variant of BC that imposes a sparsity prior on the implicitly-represented rewards.
Sparsity regularization. The issue with BC is that, when the agent encounters states that are outof-distribution with respect to Ddemo, Qθ may output arbitrary values. One solution from prior work
3Assuming a discrete action space simplifies our analysis. SQIL can be applied to continuous control tasks using existing sampling methods (Haarnoja et al., 2017; 2018), as illustrated in Section 4.3.
(Piot et al., 2014) is to regularize Qθ with a sparsity prior on the implied rewards – in particular, a penalty on the magnitude of the rewards ∑ s∈S,a∈A |Rq(s, a)| implied by Qθ via the soft Bellman equation (Equation 3), where
Rq(s, a) , Qθ(s, a)− γEs′ [ log (∑ a′∈A exp (Qθ(s ′, a′)) )] . (7)
Note that the reward function Rq is not explicitly modeled in this method. Instead, we directly minimize the magnitude of the right-hand side of Equation 7, which is equivalent to minimizing |Rq(s, a)|. The purpose of the penalty on |Rq(s, a)| is two-fold: (1) it imposes a sparsity prior motivated by prior work (Piot et al., 2013), and (2) it incorporates information about the state transition dynamics into the imitation learning objective, since Rq(s, a) is a function of an expectation over next state s′. (2) is critical for learning long-horizon behavior that imitates the demonstrations, instead of greedy maximization of the action likelihoods in standard BC. For details, see Piot et al. (2014).
Approximations for continuous states. Unlike the discrete environments tested in Piot et al. (2014), we assume the continuous state space S cannot be enumerated. Hence, we approximate the penalty ∑ s∈S,a∈A |Rq(s, a)| by estimating it from samples: transitions (s, a, s′) observed in the demonstrations Ddemo, as well as additional rollouts Dsamp periodically sampled during training using the latest imitation policy. This approximation, which follows the standard approach to constraint sampling (Calafiore & Dabbene, 2006), ensures that the penalty covers the state distribution actually encountered by the agent, instead of only the demonstrations.
To make the penalty continuously differentiable, we introduce an additional approximation: instead of penalizing the absolute value |Rq(s, a)|, we penalize the squared value (Rq(s, a))2. Note that since the reward function Rq is not explicitly modeled, but instead defined via Qθ in Equation 7, the squared penalty (Rq(s, a))2 is equivalent to the squared soft Bellman error δ2(Ddemo ∪Dsamp, 0) from Equation 1.
Regularized BC algorithm. Formally, we define the regularized BC loss function adapted from Piot et al. (2014) as
`RBC(θ) , `BC(θ) + λδ 2(Ddemo ∪ Dsamp, 0), (8)
where λ ∈ R≥0 is a constant hyperparameter, and δ2 denotes the squared soft Bellman error defined in Equation 1. The BC loss encourages Qθ to output high values for demonstrated actions at demonstrated states, and the penalty term propagates those high values to nearby states. In other words, Qθ outputs high values for actions that lead to states from which the demonstrated states are reachable. Hence, when the agent finds itself far from the demonstrated states, it takes actions that lead it back to the demonstrated states.
The RBC algorithm follows the same procedure as Algorithm 1, except that in line 4, RBC takes a gradient step on the RBC loss from Equation 8 instead of the SQIL loss.

heading: 3.3 CONNECTION BETWEEN SQIL AND REGULARIZED BEHAVIORAL CLONING

text:
The gradient of the RBC loss in Equation 8 is proportional to the gradient of the SQIL loss in line 4 of Algorithm 1, plus an additional term that penalizes the soft value of the initial state s0 (full derivation in Section A.1 of the appendix):
∇θ`RBC(θ) ∝ ∇θ ( δ2(Ddemo, 1) + λsampδ2(Dsamp, 0) + V (s0) ) . (9)
In other words, SQIL solves a similar optimization problem to RBC. The reward function in SQIL also has a clear connection to the sparsity prior in RBC: SQIL imposes the sparsity prior from RBC, by training the agent with an extremely sparse reward function – r = +1 at the demonstrations, and r = 0 everywhere else. Thus, SQIL can be motivated as a practical way to implement the ideas for regularizing BC proposed in Piot et al. (2014).
The main benefit of using SQIL instead of RBC is that SQIL is trivial to implement, since it only requires a few small changes to any standard Q-learning implementation (see Section 2). Extending SQIL to MDPs with a continuous action space is also easy, since we can simply replace Q-learning
with an off-policy actor-critic method (Haarnoja et al., 2018) (see Section 4.3). Given the difficulty of implementing deep RL algorithms correctly (Henderson et al., 2018), this flexibility makes SQIL more practical to use, since it can be built on top of existing implementations of deep RL algorithms. Furthermore, the ablation study in Section 4.4 suggests that SQIL actually performs better than RBC.

heading: 4 EXPERIMENTAL EVALUATION

text: Our experiments aim to compare SQIL to existing imitation learning methods on a variety of tasks with high-dimensional, continuous observations, such as images, and unknown dynamics. We benchmark SQIL against BC and GAIL4 on four image-based games – Car Racing, Pong, Breakout, and Space Invaders – and three state-based tasks – Humanoid, HalfCheetah, and Lunar Lander (Brockman et al., 2016; Bellemare et al., 2013; Todorov et al., 2012). We also investigate which components of SQIL contribute most to its performance via an ablation study on the Lunar Lander game. Section A.3 in the appendix contains additional experimental details.

heading: 4.1 TESTING GENERALIZATION IN IMAGE-BASED CAR RACING

text:
The goal of this experiment is to study not only how well each method can mimic the expert demonstrations, but also how well they can acquire policies that generalize to new states that are not seen in the demonstrations. To do so, we train the imitation agents in an environment with a different initial state distribution S train0 than that of the expert demonstrations Sdemo0 , allowing us to systematically control the mismatch between the distribution of states in the demonstrations and the states actually encountered by the agent. We run experiments on the Car Racing game from OpenAI Gym. To create S train0 , the car is rotated 90 degrees so that it begins perpendicular to the track, instead of parallel to the track as in Sdemo0 . This intervention presents a significant generalization challenge to the imitation learner, since the expert demonstrations do not contain any examples of states where the car is perpendicular to the road, or even significantly off the road axis. The agent must learn to make a tight turn to get back on the road, then stabilize its orientation so that it is parallel to the road, and only then proceed forward to mimic the expert demonstrations.
The results in Figure 1 show that SQIL and BC perform equally well when there is no variation in the initial state. The task is easy enough that even BC achieves a high reward. Note that, in the unperturbed condition (right column), BC substantially outperforms GAIL, despite the wellknown shortcomings of BC. This indicates that the adversarial optimization in GAIL can substantially hinder
learning, even in settings where standard BC is sufficient. SQIL performs much better than BC when starting from S train0 , showing that SQIL is capable of generalizing to a new initial state distribution, while BC is not. SQIL learns to make a tight turn that takes the car through the grass and back onto the road, then stabilizes the car’s orientation so that it is parallel to the track, and then proceeds forward like the expert does in the demonstrations. BC tends to drive straight ahead into the grass instead of turning back onto the road.
4 For all the image-based tasks, we implement a version of GAIL that uses deep Q-learning (GAIL-DQL) instead of TRPO as in the original GAIL paper (Ho & Ermon, 2016), since Q-learning performs better than TRPO in these environments, and because this allows for a head-to-head comparison of SQIL and GAIL: both algorithms use the same underlying RL algorithm, but provide the agent with different rewards – SQIL provides constant rewards, while GAIL provides learned rewards. We use the standard GAIL-TRPO method as a baseline for all the low-dimensional tasks, since TRPO performs better than Q-learning in these environments.
The original GAIL method implicitly encodes prior knowledge – namely, that terminating an episode is either always desirable or always undesirable. As pointed out in Kostrikov et al. (2019), this makes comparisons to alternative methods unfair. We implement the unbiased version of GAIL proposed by Kostrikov et al. (2019), and use this in all of the experiments. Comparisons to the biased version with implicit termination knowledge are included in Section A.2 in the appendix.
SQIL outperforms GAIL in both conditions. Since SQIL and GAIL both use deep Q-learning for RL in this experiment, the gap between them may be attributed to the difference in the reward functions they use to train the agent. SQIL benefits from providing a constant reward that does not require fitting a discriminator, while GAIL struggles to train a discriminator to provide learned rewards directly from images.

heading: 4.2 IMAGE-BASED EXPERIMENTS ON ATARI

text:
The results in Figure 2 show that SQIL outperforms BC on Pong, Breakout, and Space Invaders – additional evidence that BC suffers from compounding errors, while SQIL does not. SQIL also outperforms GAIL on all three games, illustrating the difficulty of using GAIL to train an imagebased discriminator, as in Section 4.1.
4.3 INSTANTIATING SQIL FOR CONTINUOUS CONTROL IN LOW-DIMENSIONAL MUJOCO
The experiments in the previous sections evaluate SQIL on MDPs with a discrete action space. This section illustrates how SQIL can be adapted to continuous actions. We instantiate SQIL using soft actor-critic (SAC) – an off-policy RL algorithm that can solve continuous control tasks (Haarnoja et al., 2018). In particular, SAC is modified in the following ways: (1) the agent’s experience replay buffer is initially filled with expert demonstrations, where rewards are set to r = +1, (2) when taking gradient steps to fit the agent’s soft Q function, a balanced number of demonstration experiences and new experiences (50% each) are sampled from the replay buffer, and (3) the agent observes rewards of r = 0 during its interactions with the environment, instead of an extrinsic reward signal that specifies the desired task. This instantiation of SQIL is compared to GAIL on the Humanoid (17 DoF) and HalfCheetah (6 DoF) tasks from MuJoCo.
The results show that SQIL outperforms BC and performs comparably to GAIL on both tasks, demonstrating that SQIL can be successfully deployed on problems with continuous actions, and that SQIL can perform well even with a small number of demonstrations. This experiment also illustrates how SQIL can be run on top of SAC or any other off-policy value-based RL algorithm.

heading: 4.4 ABLATION STUDY ON LOW-DIMENSIONAL LUNAR LANDER

text:
We hypothesize that SQIL works well because it combines information about the expert’s policy from demonstrations with information about the environment dynamics from rollouts of the imitation policy periodically sampled during training. We also expect RBC to perform comparably to SQIL, since their objectives are similar. To test these hypotheses, we conduct an ablation study using the Lunar Lander game from OpenAI Gym. As in Section 4.1, we control the mismatch between the
distribution of states in the demonstrations and the states encountered by the agent by manipulating the initial state distribution. To create S train0 , the agent is placed in a starting position never visited in the demonstrations.
In the first variant of SQIL, λsamp is set to zero, to prevent SQIL from using additional samples drawn from the environment (see line 4 of Algorithm 1). This comparison tests if SQIL really needs to interact with the environment, or if it can rely solely on the demonstrations. In the second condition, γ is set to zero to prevent SQIL from accessing information about state transitions (see Equation 1 and line 4 of Algorithm 1). This comparison tests if SQIL is actually extracting information about the dynamics from the samples, or if it can perform just as well with a naı̈ve regularizer (setting γ to zero effectively imposes a penalty on the L2-norm of the soft Q values instead of the squared soft Bellman error). In the third condition, a uniform random policy is used to sample additional rollouts, instead of the imitation policy πθ (see line 6 of Algorithm 1). This comparison tests how important it is that the samples cover the states encountered by the agent during training. In the fourth condition, we use RBC to optimize the loss in Equation 8. instead of using SQIL to optimize the loss in line 4 of Algorithm 1. This comparison tests the effect of the additional V (s0) term in RBC vs. SQIL (see Equation 9).
The results in Figure 4 show that all methods perform well when there is no variation in the initial state. When the initial state is varied, SQIL performs significantly better than BC, GAIL, and the ablated variants of SQIL. This confirms our hypothesis that SQIL needs to sample from the environment using the imitation policy, and relies on information about the dynamics encoded in the samples. Surprisingly, SQIL outperforms RBC by a large margin, suggesting that the penalty on the soft value of the initial
state V (s0), which is present in RBC but not in SQIL (see Equation 9), degrades performance.

heading: 5 DISCUSSION AND RELATED WORK

text:
Related work. Concurrently with SQIL, two other imitation learning algorithms that use constant rewards instead of a learned reward function were developed (Sasaki et al., 2019; Wang et al., 2019). We see our paper as contributing additional evidence to support this core idea, rather than proposing a competing method. First, SQIL is derived from sparsity-regularized BC, while the prior methods are derived from an alternative formulation of the IRL objective (Sasaki et al., 2019) and from support estimation methods (Wang et al., 2019), showing that different theoretical approaches independently lead to using RL with constant rewards as an alternative to adversarial training – a sign that this idea may be a promising direction for future work. Second, SQIL is shown to outperform BC and GAIL in domains that were not evaluated in Sasaki et al. (2019) or Wang et al. (2019) – in particular, tasks with image observations and significant shift in the state distribution between the demonstrations and the training environment.
Summary. We contribute the SQIL algorithm: a general method for learning to imitate an expert given action demonstrations and access to the environment. Simulation experiments on tasks with high-dimensional, continuous observations and unknown dynamics show that our method outperforms BC and achieves competitive results compared to GAIL, while being simple to implement on top of existing off-policy RL algorithms.
Limitations and future work. We have not yet proven that SQIL matches the expert’s state occupancy measure in the limit of infinite demonstrations. One direction for future work would be to rigorously show whether or not SQIL has this property. Another direction would be to extend SQIL to recover not just the expert’s policy, but also their reward function; e.g., by using a parameterized reward function to model rewards in the soft Bellman error terms, instead of using constant rewards. This could provide a simpler alternative to existing adversarial IRL algorithms.

heading: A APPENDIX

text:
A.1 DERIVATION OF RBC GRADIENT
Let τ = (s0, a0, s1, ..., sT ) denote a rollout, where sT is an absorbing state. Let V denote the soft value function,
V (s) , log (∑ a∈A exp (Qθ(s, a)) ) . (10)
Splitting up the squared soft Bellman error terms for Ddemo and Dsamp in Equation 8,
∇`RBC(θ) = ∑
τ∈Ddemo
T−1∑ t=0 −(∇Qθ(st, at)−∇V (st))
+ λdemo ∑
τ∈Ddemo
T−1∑ t=0 ∇(Qθ(st, at)− γV (st+1))2 + λsamp∇δ2(Dsamp, 0)
= ∑
τ∈Ddemo
T−1∑ t=0 ∇V (st)− γ∇V (st+1)
+ λdemo∇δ2 ( Ddemo, 1
2λdemo
) + λsamp∇δ2(Dsamp, 0). (11)
Setting γ , 1 turns the inner sum in the first term into a telescoping sum:
(11) = ∑
τ∈Ddemo
(∇V (s0)−∇V (sT )) + λdemo∇δ2 ( Ddemo, 1
2λdemo
) + λsamp∇δ2(Dsamp, 0). (12)
Since sT is assumed to be absorbing, V (sT ) is zero. Thus,
(12) = ∑
s0∈Ddemo
∇V (s0) + λdemo∇δ2 ( Ddemo, 1
2λdemo
) + λsamp∇δ2(Dsamp, 0), (13)
In our experiments, we have that all the demonstration rollouts start at the same initial state s0.5 Thus,
(13) ∝ ∇ ( δ2(Ddemo, 1) + λsampδ2(Dsamp, 0) + V (s0) ) , (14)
where λsamp ∈ R≥0 is a constant hyperparameter.
A.2 COMPARING THE BIASED AND UNBIASED VARIANTS OF GAIL
As discussed in Section 4, to correct the original GAIL method’s biased handling of rewards at absorbing states, we implement the suggested changes to GAIL in Section 4.2 of Kostrikov et al. (2019): adding a transition to an absorbing state and a self-loop at the absorbing state to the end of each rollout sampled from the environment, and adding a binary feature to the observations indicating whether or not a state is absorbing. This enables GAIL to learn a non-zero reward for absorbing states. We refer to the original, biased GAIL method as GAIL-DQL-B and GAIL-TRPO-B, and the unbiased version as GAIL-DQL-U and GAIL-TRPO-U.
The mechanism for learning terminal rewards proposed in Kostrikov et al. (2019) does not apply to SQIL, since SQIL does not learn a reward function. SQIL implicitly assumes a reward of zero at absorbing states in demonstrations. This is the case in all our experiments, which include some environments where terminating the episode is always undesirable (e.g., walking without falling down) and other environments where success requires terminating the episode (e.g., landing at a target), suggesting that SQIL is not sensitive to the choice of termination reward, and neither significantly benefits nor is significantly harmed by setting the termination reward to zero.
5Demonstration rollouts may still vary due to the stochasticity of the expert policy.
Car Racing. The results in Figure 5 show that both the biased (GAIL-DQL-B) and unbiased (GAILDQL-U) versions of GAIL perform equally poorly. The problem of training an image-based discriminator for this task may be difficult enough that even with an unfair bias toward avoiding crashes that terminate the episode, GAIL-DQL-B does not perform better than GAIL-DQL-U.
Atari. The results in Figure 6 show that SQIL outperforms both variants of GAIL on Pong and the unbiased version of GAIL (GAIL-DQL-U) on Breakout and Space Invaders, but performs comparably to the biased version of GAIL (GAIL-DQL-B) on Space Invaders and worse than it on Breakout. This may be due to the fact that in Breakout and Space Invaders, the agent has multiple lives – five in Breakout, and three in Space Invaders – and receives a termination signal that the episode has ended after losing each life. Thus, the agent experiences many more episode terminations than in Pong, exacerbating the bias in the way the original GAIL method handles rewards at absorbing states. Our implementation of GAIL-DQL-B in this experiment provides a learned reward of r(s, a) = − log (1−D(s, a)), where D is the discriminator (see Section A.3 in the appendix for details). The learned reward is always positive, while the implicit reward at an absorbing state is zero. Thus, the agent is inadvertently encouraged to avoid terminating the episode. For Breakout and Space Invaders, this just happens to be the right incentive, since the objective is to stay alive as long as possible. GAIL-DQL-B outperforms SQIL in Breakout and performs comparably to SQIL in Space Invaders because GAIL-DQL-B is accidentally biased in the right way.
Lunar Lander. The results in Figure 7 show that when the initial state is varied, SQIL outperforms the unbiased variant of GAIL (GAIL-TRPO-U), but underperforms against the biased version of GAIL (GAIL-TRPO-B). The latter result is likely due to the fact that the implementation of GAILTRPO-B we used in this experiment provides a learned reward of r(s, a) = log (D(s, a)), where D is the discriminator (see Section A.3 in the appendix for details). The learned reward is always negative, while the implicit reward at an absorbing state is zero. Thus, the agent is inadvertently encouraged to terminate the episode quickly. For the Lunar Lander game, this just happens to be the right incentive, since the objective is to land on the ground and thereby terminate the episode. As in the Atari experiments, GAIL-TRPO-B performs better than SQIL in this experiment because GAIL-TRPO-B is accidentally biased in the right way.
Humanoid-v1
A.3 IMPLEMENTATION DETAILS
To ensure fair comparisons, the same network architectures were used to evaluate SQIL, GAIL, and BC. For Lunar Lander, we used a network architecture with two fully-connected layers containing 128 hidden units each to represent the Q network in SQIL, the policy and discriminator networks in GAIL, and the policy network in BC. For Car Racing, we used four convolutional layers (following (Ha & Schmidhuber, 2018)) and two fully-connected layers containing 256 hidden units each. For Humanoid and HalfCheetah, we used two fully-connected layers containing 256 hidden units each. For Atari, we used the convolutional neural network described in (Mnih et al., 2015) to represent the Q network in SQIL, as well as the Q network and discriminator network in GAIL.
To ensure fair comparisons, the same demonstration data were used to train SQIL, GAIL, and BC. For Lunar Lander, we collected 100 demonstration rollouts. For Car Racing, Pong, Breakout, and Space Invaders, we collected 20 demonstration rollouts. Expert demonstrations were generated from scratch for Lunar Lander using DQN (Mnih et al., 2015), and collected from open-source pretrained policies for Car Racing (Ha & Schmidhuber, 2018) as well as Humanoid and HalfCheetah (Dhariwal et al., 2017). The Humanoid demonstrations were generated by a stochastic expert policy, while the HalfCheetah demonstrations were generated by a deterministic expert policy; both experts were trained using TRPO.6 We used two open-source implementations of GAIL: (Fu et al., 2017) for Lunar Lander, and (Dhariwal et al., 2017) for MuJoCo. We adapted the OpenAI Baselines implementation of GAIL to use soft Q-learning for Car Racing and Atari. Expert demonstrations were generated from scratch for Atari using DQN.
For Lunar Lander, we set λsamp = 10−6. For Car Racing, we set λsamp = 0.01. For all other environments, we set λsamp = 1.
6https://drive.google.com/drive/folders/1h3H4AY_ZBx08hz-Ct0Nxxus-V1melu1U
SQIL was not pre-trained in any of the experiments. GAIL was pre-trained using BC for HalfCheetah, but was not pre-trained in any other experiments.
In standard implementations of soft Q-learning and SAC, the agent’s experience replay buffer typically has a fixed size, and once the buffer is full, old experiences are deleted to make room for new experiences. In SQIL, we never delete demonstration experiences from the replay buffer, but otherwise follow the standard implementation.
We use Adam (Kingma & Ba, 2014) to take the gradient step in line 4 of Algorithm 1.
The BC and GAIL performance metrics in Section 4.3 are taken from (Dhariwal et al., 2017).7
The GAIL and SQIL policies in Section 4.3 are set to be deterministic during the evaluation rollouts used to measure performance.
7https://github.com/openai/baselines/blob/master/baselines/gail/result/ gail-result.md