Bayesian Disturbance Injection:
Robust Imitation Learning of Flexible Policies

A key objective of imitation learning is to ensure that models can capture the variation and stochasticity inherent in human motion. Classical dynamical frameworks for learning trajectories from demonstrations include Dynamic Movement Primitives (DMPs), which can generalize the learned trajectories to new situations (i.e., goal location or speed). However, this generalization depends on the heuristic, thus unsuitable for learning state-dependent feedback policies [7, 8, 9].

Gaussian Mixture Regression (GMR) is an non-parametric, intuitive means to learn trajectories or policies from demonstrators in the state-action-space. In this, Gaussian Mixture Modelling (GMM) [10] is used as a basis function to capture non-linearities during learning, and has been utilized in imitation learning that deals with human demonstrations [11]. However, GMR requires to engineer features (e.g., Gaussian initial conditions) by hand to deal with high-dimensional systems. [12].

As an alternative, Gaussian process regression (GPR) deals with implicit (high-dimensional) feature spaces with kernel functions. It thus can directly deal with high-dimensional observations without explicitly learning in this high-dimensional space [13]. In particular, Overlapping Mixtures of Gaussian Processes (OMGP) [14] learns a multi-modal distribution by overlapping multiple GPs, and has been employed as a policy model with multiple optimal actions on flexible task learning of robotic policies [15]. To further reduce a priori tuning, Infinite Overlapping Mixtures of Gaussian Processes (IOMGP) [6] requires only an upper bound of the number of GPs to be estimated. As such, IOMGP is an intuitive means of learning complex multi-modal policies from unlabeled human demonstration data and is employed in this paper.

While flexibility is key to capturing demonstrator motion, a major issue limiting application of learned policies is the problem of covariate shift. Specifically, environment variations (e.g., manipulator starting position) induces differences between the policy distribution as learned by the manipulator and the actual task distribution during application.

A more general approach to minimizing the covariate shift in imitation learning, is Dataset Aggregation (DAgger) [16], whereby when the robot moves to a state not included in the training data, the expert teaches the optimal recovery. However, this approach has limited applicability in practice due to the risk of running a poorly learned robot and the tediousness to have human experts continue to teach the robot the optimal actions.

An intuitive approach to robustifying learned policies against sources of error, without needing to a priori specify task-relevant learning parameters, is to exploit phenomenon similar to persistence excitation [17]. In this, noise is injected into the expert’s demonstrated actions, and the recovery behavior of the expert is learned from this perturbation. In an imitation learning context, Disturbances for Augmenting Robot Trajectories (DART) [18] exploits this phenomenon for learning a deterministic policy model with a single optimal action. Additionally, DART is well suited to creating a richer dataset, by concurrently determining the optimal noise level to be injected into the demonstrated actions during policy learning.

However, the algorithm proposed to achieve DART [18] has a major limitation. It assumes a uni-model deterministic policy, which is unsuitable for many real-world imitation learning tasks, that may consist of multiple optimal actions.

In contrast to that approach, this paper explores a novel method that can eliminate such a severe limitation by incorporating a prior distribution on policy then optimizing both the policy and the injection noise parameters simultaneously via variational Bayesian learning. As such, this novel imitation learning framework improves robustness while maintaining flexibility.

The objective of imitation learning is to learn a control policy by imitating the action from the expert’s demonstration data. A dynamics model is denoted as Markovian with a state $\mathbf{s}_{t}\in\mathbb{R}^{Q}$, an action $a_{t}\in\mathbb{R}$ and a state transition distribution $p(\mathbf{s}_{t+1}\mid\mathbf{s}_{t},a_{t})$. A policy $\pi(a_{t}\mid\mathbf{s}_{t})$ decides an action from a state. A trajectory $\tau=(\mathbf{s}_{0},{a}_{0},\mathbf{s}_{1},{a}_{1}\dots a_{T-1},\mathbf{s}_{T})$ which is a sequence of state-action pairs of $T$ steps. The trajectory distribution is indicated as:

A key aspect of imitation learning is to replicate the expert’s behavior, and as such the function which computes the similarity of two policies using trajectories is defined as:

A learned policy $\pi^{R}$ is obtained by solving the following optimization problem using a trajectory collected by an expert’s policy $\pi^{*}$:

In imitation learning, a learned policy may suffer from the problem of error compounding, caused by covariate shift. This is defined as the distributional difference between the trajectories in data collection, and those in testing:

To learn policies that are robust to covariate shift, DART has previously been proposed [18] for imitation learning problems. In this, expert demonstrations are injected with noise to produce a richer set of demonstrated actions. The level of injection noise is optimized iteratively to reduce covariate shift during data collection.

In this, it is assumed that the injection noise is sampled from a Gaussian distribution as $\epsilon_{t}\sim\mathcal{N}(0,\sigma^{2}_{k})$, where $k$ is the number of optimizations. The injection noise $\epsilon_{t}$ is added to the expert’s action $a^{*}_{t}$.

The noise injected expert’s trajectory distribution is denoted as $p(\tau\mid\pi^{*},\sigma^{2}_{k})$ and the trajectory distribution with learned policy as $p(\tau\mid\pi^{R}_{k})$. To reduce covariate shift, DART introduces upper bound of covariate shift by Pinsker’s inequality as:

where, $\mathrm{KL}(\cdot\mid\mid\cdot)$ is Kullback-Leibler divergence. However, the upper bound is intractable since the learned policy’s trajectory distribution $p(\tau\mid\pi^{R}_{k})$ is unknown. Therefore, DART solves the upper bound using the noise injected expert’s trajectory distribution instead of the learned policy’s trajectory distribution. Optimal injection noise distribution is optimized as:

The optimized injection noise distribution can be interpreted as the likelihood of the learned policy.

The learned policy $\pi^{R}_{k}$ used in the optimization of $\sigma^{2}_{k+1}$ is obtained by following:

To learn a multi-modal policy, the policy prior is considered as the product of infinite GPs, inspired by IOMGP. Fig. 2 shows a policy model in which expert’s actions $\mathbf{a^{*}}=[a^{*}_{n}]_{n=1}^{N}$ are estimated by $\mathbf{f}^{(m)},\mathbf{Z},\{\sigma_{k+1}^{2}\}$ where $N=\sum_{j=1}^{k}N_{j}$ , $N_{j}$ is a size of the dataset that collected at $j$-th iteration. The latent function $\mathbf{f}^{(m)}$ is the output of $m$-th GP given state $\mathbf{S}=[\mathbf{s}_{n}]_{n=1}^{N}$. To allocate $n$-th expert’s action $a_{n}^{*}$ to $m$-th latent function $\mathbf{f}^{(m)}$, indicator matrix $\mathbf{Z}\in\mathbb{R}^{N\times\infty}$ is defined. To estimate the optimal number of GPs, a random variable $v_{m}$ quantifies the uncertainty assigned to $\mathbf{f}^{(m)}$. In addition, the set of injection noise parameters, $\{\sigma_{k+1}^{2}\}=\{\sigma_{j+1}^{2}\}_{j=1}^{k}$, indicate the distribution of the expert’s action $a_{n}^{*}$ in the distance from the latent function $\mathbf{f}^{(m)}$.

The set of latent functions is denoted as $\{{\mathbf{f}}^{(m)}\}=\{\mathbf{f}^{(m)}\}^{\infty}_{m=1}$ and a GP prior is given by :

where $\mathbf{K}^{(m)}=\mathrm{k}^{(m)}({\mathbf{S}},{\mathbf{S}})$ is the $m$-th kernel gram matrix with the kernel function $\mathrm{k}^{(m)}(\cdot,\cdot)$ and a kernel hyperparameter $\omega^{(m)}$. Let $\{\boldsymbol{\omega}^{(m)}\}=\{\omega^{(m)}\}^{\infty}_{m=1}$ be the set of hyperparameters of infinite number of kernel functions.

Then the Stick Breaking Process (SBP) [19] is used as a prior of $\mathbf{Z}$, which can be interpreted as an infinite mixture model as follows:

Note that the implementation of variational Bayesian learning cannot deal with infinite-dimensional vectors, so the $\infty$ component is replaced with a predefined upper bound of the mixtures $M$. ${{v}_{m}}$ is a random variable indicating the probability that the data corresponds to the $m$-th GP. Thus, it is possible to estimate the optimal number of GPs with a high probability of allocation starting from an infinite number of GPs. $\beta$ is a hyperparameter of SBP denoting the level of concentration of the data in the cluster.

The above policy model differs from the IOMGP model for regression [6]; our model employs the set of injection noise parameters $\{{\sigma}_{k+1}^{2}\}$ because of the injection noise distribution is updated at each iteration. The next injection noise optimized by (III-B) deals with the training data collected in the current iteration, while the policy inference (7) deals with all the training data gathered up to the current iteration. Thus, the heteroscedastic Gaussian noise is defined as:

showing the association between the injection noise parameters and the training data collected for each iteration. In this, $\mathbf{1}_{N_{i}}$ is a vector whose size is $N_{i}$ and all components are one. In addition, the value of the likelihood (III-B) does not change even if the mean and the input are replaced, the likelihood for the $\{\mathbf{f}^{(m)}\},\mathbf{Z},\{{\sigma}_{k+1}^{2}\}$ is derived as :

This formulation is described in a graphical model that defines the relationship between the variables as shown in Fig. 2, and the joint distribution of the model as :

where $\Omega=(\{\boldsymbol{\omega}^{(m)}\},\{{\sigma}_{k+1}^{2}\},\beta)$ represents a set of hyperparameters.

Bayesian learning is a framework that estimates the posterior distributions of the policies and their predictive distributions for new input data rather than point estimates the policy parameters. To obtain the posterior and predictive distributions, the marginal likelihood is calculated as :

However, it is intractable to calculate the log marginal likelihood of (IV-B) analytically. Therefore, the variational lower bound is derived as the objective function of variational learning. The true posterior distribution is approximated by the variational posterior distribution, which maximizes the variational lower bound. As a common fashion of variational inference, the variational posterior distribution is assumed to be factorized among all latent variables (known as the mean-field approximation [20]) as follows:

The variational lower bound $\mathcal{L}(q,\Omega)$ is derived from this assumption by applying the Jensen inequality to the log marginal likelihood, as:

In addition, the optimization formulation is derived using the Expectation-Maximization (EM) algorithm. The variational posterior distribution $q$ is optimized with a fixed hyperparameter $\Omega$ in E-step, and the hyperparameter $\Omega$ is optimized with a fixed variational posterior distribution $q$ in M-step with:

See the SectionVII-A,B for details of $q$ update laws and lower bound of marginal likelihood. And a summary of the proposed method is shown in Algorithm 1.

Using the hyper-parameter $\Omega$ and the variational posterior distribution $q$, optimized by variational Bayesian learning, the predictive distribution of the $m$-th action $a_{*}^{(m)}$ and variance $\sigma_{*}^{2(m)}$ on an current state $\mathbf{s}_{*}$ can be computed analytically as (see [13]).

An initial experiment is presented to characterise the flexibility and robustness of the proposed approach, for learning and completing tasks in an environment. Specifically, a standard manipulator learning task (table-sweeping) involving multi-object handling is performed in the V-REP [22] environment, as shown in Fig. 3-(a).

In this experiment, learned policies are evaluated in terms of ability to flexibly learn tasks with multiple optimal actions (e.g., the order in which to sweep the objects from the table), and well as robustness to environmental covariance shift inducing disturbances (e.g., friction between the objects and table inducing variations of object movement).

Update of $q(\mathbf{f}^{(m)})$:

Update of $q(\mathbf{Z})$:

where, $\psi(\cdot)$ is the digamma function.

Update of $q(v_{m})$: