Generalization Guarantees for Imitation Learning

Multi-modal imitation learning. Imitation learning is commonly used in manipulation [13, 14, 2] and navigation tasks [15, 16] to accelerate training by injecting expert knowledge. Often, imitation data can be multi-modal, e.g., an expert may choose to grasp anywhere along the rim of a mug. Recently, a large body of work uses latent variable models to capture such multi-modality in robotic settings [17, 18, 19, 20, 10]. While these papers use multi-modal data to diversify imitated behavior, we embed the multi-modality into the prior policy distribution to accelerate the “fine-tuning” and produce better generalization bounds and empirical performance (see Sec. 4.3 and A4).

Learning from imperfect demonstrations. Another challenge with imitation learning is that expert demonstrations can often be imperfect [21]. Moreover, policies trained using only off-policy data can also cause cascading errors [22] when the robot encounters unseen states. One approach to addressing these challenges is to “fine-tune” an imitation-learned policy to improve empirical performance in novel environments [21, 23]. Other techniques that mitigate these issues include collecting demonstrations with noisy dynamics [24], augmenting observation-action data of demonstrations using noise [13], and training using a hybrid reward for both imitation and task success [2]. Some recent work also explores generalization in longer horizon tasks [23, 25]. However, none of these techniques provide rigorous generalization guarantees for imitated policies deployed in novel environments for robotic systems with discrete/continuous state and action spaces, nonlinear/hybrid dynamics, and rich sensing (e.g., vision). This is the primary focus of this paper.

Generalization guarantees for learning-based control. The PAC-Bayes Control framework [26, 27, 28] provides a way to make safety and generalization guarantees for learning-based control. In this work, we extend this framework to make guarantees on policies learned via imitation learning. To this end, we propose a two-phase training pipeline for learning a prior distribution over policies via imitation learning and then fine-tuning this prior by optimizing a PAC-Bayes generalization bound. By leveraging multi-modal expert demonstrations, we are able to obtain significantly stronger generalization guarantees than [27, 28], which either employ simple heuristics to choose the prior or train the prior from scratch.

The goal of the first training stage is to obtain a prior policy distribution $P_{0}$ by cloning expert demonstrations $\zeta_{i}$. Behavioral cloning is a straightforward strategy to make robots mimic expert behavior. While simple discriminative models fail to capture diverse expert behavior, generative models such as variational autoencoders (VAEs) can embed such multi-modality in latent variables [11, 18]. We further use a conditional VAE (cVAE) [11] to condition the action output $a$ on both the latent $z$ and observation $o$ (Fig. 2). The latent $z$ encodes both $o$ and $a$ from expert demonstrations $\{\zeta_{i}\}_{i=1}^{M}$. In the example of grasping mugs, intuitively we can consider that $z$ encodes the mug center location and the relative-to-center grasp pose from depth images and demonstrated grasps. Both pieces of information are necessary for generating successful, diverse grasps along the rim using different sampled $z$ (Fig. 4(a)).

While grasping mugs can be achieved by executing an open-loop grasp from above, other tasks such as pushing boxes and navigating indoor environments require continuous, closed-loop actions. We embed either a short sequence of observation/action pairs ($T=3$ steps) or the entire trajectory (as many as 50 steps) into a single latent $z$. Thus a sampled latent state can represent local, reactive actions (e.g. turn left to avoid a chair during navigation), or global “plans” (e.g. push at the corner of an object throughout the horizon while manipulating it).

Fig. 2 shows the graphical model of the cVAE in our work. The encoder $q_{\phi}(z|o,a)$, parameterized by weights $\phi$ of a neural network, samples a latent variable $z$ conditioned by each demonstration $\zeta_{i}:=\{(o_{t},a_{t})\}_{t=1}^{T}$. The decoder, $\pi_{\theta,z}:o\mapsto a$, parameterized by the weights $\theta$ of another neural network and the sampled latent variable $z$, reconstructs the action from the observation at each step. Details of the neural network architecture are provided in A5. In the loss function below, the distribution $p(z)$ is chosen as a multivariate unit Gaussian. A KL regularization loss constrains the conditional distribution of $z$ to be close to $p(z)$. Let $\mathcal{L}_{\rm rec}(\pi_{\scaleto{\theta}{4pt},z}(o_{t}),a_{t})$ be the reconstruction loss between the predicted action and the expert’s. The parameter $\lambda>0$ balances the two losses:

The primary outcomes of behavioral cloning are: (i) a latent distribution $p(z)=\mathcal{N}(0,I)$ and (ii) weights $\theta^{*}$ of the decoder network $\pi_{\scaleto{\theta}{4pt},z}$ that together encode multi-modal policies from experts. We now restrict the weights $\theta$ of the decoder network $\pi_{\scaleto{\theta}{4pt},z}$ to $\theta^{*}$ giving rise to the space of policies $\Pi:=\{\pi_{\scaleto{\theta^{*}}{4pt},z}:\mathcal{O}\to\mathcal{A}~{}|~{}z\in\mathbb{R}^{n_{z}}\}$ parameterized by $z$. Hence, the latent distribution $p(z)$ can be equivalently viewed as a distribution on the space $\Pi$ of policies. In the next section, we will consider $p(z)$ as a prior distribution $P_{0}$ on $\Pi$ and “fine-tune” it by searching for a posterior distribution $P$ in the space $\mathcal{P}$ of probability distributions on $\Pi$ by solving (2). In particular, we choose $\mathcal{P}$ as the space of Gaussian probability distributions with diagonal covariance $\mathcal{N}(\mu,\Sigma)$. For the sake of notational convenience, let $\sigma\in\mathbb{R}^{d}$ be the element-wise square-root of the diagonal of $\Sigma$, and define $\psi:=(\mu,\sigma)$, $P=\mathcal{N}_{\psi}:=\mathcal{N}(\mu,\text{diag}(\sigma^{2}))$, and $P_{0}=\mathcal{N}_{\psi_{0}}:=\mathcal{N}(0,I)$.

Although behavioral cloning provides a meaningful policy distribution $P_{0}$, the policies drawn from this distribution can fail when deployed in novel environments due to: unsafe or non-generalizable demonstrations by the expert, or the inability of the cVAE training to accurately infer the expert’s policies. In this section, we leverage the PAC-Bayes Control framework introduced in [26, 27] to “fine-tune” the prior policy distribution $P_{0}$ and provide “certificates” of generalization for the resulting posterior policy distribution $P$. In particular, we will tune the distribution by approximately minimizing the true expected cost $C_{\mathcal{D}}(P)$ in equation (2), thus promoting generalization to environments drawn from $\mathcal{D}$ that are different from $S$. Although $C_{\mathcal{D}}(P)$ cannot be computed due to the lack of an explicit characterization of $\mathcal{D}$, the PAC-Bayes framework allows us to obtain an upper bound $C_{\textup{PAC}}$ for $C_{\mathcal{D}}(P)$, which can be computed despite our lack of knowledge of $\mathcal{D}$; see Theorem 1.

To introduce the PAC-Bayes generalization bounds, we will first define the empirical cost of $P$ as the average expected cost across training environments in $S$:

The following theorem can then be used to bound the true expected cost $C_{\mathcal{D}}(P)$ from Sec. 2.

Intuitively, minimizing the upper bound $C_{\textup{PAC}}$ can be viewed as minimizing the empirical cost $C_{S}(P)$ along with a regularizer $R$ that prevents overfitting by penalizing the deviation of the posterior from the prior. Due to the presence of a “blackbox” physics simulator for rollouts, the gradient of $C_{S}(P)$ cannot be computed analytically. Thus we employ blackbox optimizers for minimizing $C_{\textup{PAC}}$.

In practice we use Natural Evolutionary Strategies (NES) [31] that transforms the ES gradient to the natural gradient [32] to accelerate training. During each epoch, for each of the $N$ environments we sample a certain number of $z$’s from the posterior $\mathcal{N}_{\psi}$ and then compute the corresponding empirical costs by performing rollouts in simulation. Sampled $z$’s and their empirical costs $C(\pi_{\scaleto{\theta^{*}}{4pt},z};E)$ are then used in (5) and (6) to compute the gradient estimate, which is passed to the Adam optimizer [33] to update $\psi$. The training loop is visualized in Fig. 3.

Computing the final bound. After ES training, we can calculate the generalization bound using the optimal $\psi^{*}$. First, note that the empirical cost $C_{S}(P)=C_{S}(\mathcal{N}_{\psi^{*}})$ involves an expectation over the posterior and thus cannot be computed in closed form. Instead, it can be estimated by sampling a large number of policies $z_{1},\cdots,z_{L}$ from $\mathcal{N}_{\psi^{*}}:\hat{C}_{S}(\mathcal{N}_{\psi}^{*}):=\frac{1}{NL}\sum_{E\in S}\sum_{i=1}^{L}C(\pi_{\scaleto{\theta^{*}}{4pt},z_{i}};E)$, and the error due to finite sampling can be bounded using a sample convergence bound $\bar{C}_{S}$ [34]. The final bound $C_{\textup{bound}}(\mathcal{N}_{\psi}^{*})\geq C_{\mathcal{D}}(\mathcal{N}_{\psi}^{*})$ is obtained from $\bar{C}_{S}$ and $R(\mathcal{N}_{\psi}^{*},P_{0})$ by a slight tightening of $C_{\textup{PAC}}$ from Theorem 1 using the KL-inverse function [27]. Please refer to A1, A2, A3 for detailed derivations and implementations.

Overall, our approach provides generalization guarantees in novel environments for policies learned from imitation learning: as policies are randomly sampled from the posterior $\mathcal{N}_{\psi}^{*}$ and applied in test environments, the expected success rate over all test environments is guaranteed to be at least $1-C_{\textup{bound}}(\mathcal{N}_{\psi}^{*})$ (with probability $1-\delta$ over the sampling of training environments; $\delta=0.01$ for all examples in Sec. 4).

The goal is to grasp and lift a mug from the table using a Franka Panda Arm (Fig. 5). An open-loop action $a_{\textup{grasp}}=(x,y,z,\theta)$ is applied for each rollout and corresponds to desired 3D positions and yaw orientation of the grasp. The action is computed based on an observation of a $128\times 128$ depth image from an overhead camera.

We gathered 50 mugs of diverse shapes and dimensions from the ShapeNet dataset [36]. These mugs are split into 3 sets for expert demonstrations, PAC-Bayes Control training environments $S$, and test environments $S^{\prime}$. They are then randomly scaled in all dimensions into multiple different mugs. Their masses are sampled from a uniform distribution. Each training or test environment consists of a unique mug from the set and a unique initial SE(2) pose on the table. A rollout is considered successful (zero cost) if the center of mass (COM) of the mug is lifted by 10 cm and the gripper palm makes no contact with the mug; otherwise, a cost of 1 is assigned to the rollout.

Expert data. In each of 60 environments, we specify 5 grasp poses along the rim. The initial depth image of the scene and corresponding grasp poses are recorded. It takes about an hour to collect the 300 trials.

Prior performance. Each pair of initial observation and action from expert data is embedded into a latent $z\in\mathbb{R}^{10}$. Thus the length of time sequence in each demonstration $\zeta_{i}$ is $T_{\textup{grasp}}=1$. The reconstruction loss $\mathcal{L}_{\textup{rec,grasp}}$ is a combination of $l_{2}$ and $l_{1}$ loss between predicted and expert’s actions. The prior policy distribution achieves 83.3% success in novel environments in simulation. Shown in Fig. 4(a), the prior $P_{0}$ captures the multi-modality of expert data: different latent $z$ sampled from $P_{0}$ generate a diverse set of grasps along the rim.

Posterior performance. $N=500$ environments are used for fine-tuning via PAC-Bayes. The resulting PAC-Bayes bound $C_{\textup{bound}}^{*}:=C_{\textup{bound}}(P)$ of the posterior $P$ is 0.070. Thus, with probability 0.99 the optimized posterior policy is guaranteed to have an expected success rate of 93.0% in novel environments (assuming that they are drawn from the same underlying distribution $\mathcal{D}$ as the training examples). The policy is then evaluated on 500 test environments in simulation and the success rate is 98.4%. Fig. 4(b) shows sampled grasps of a mug using the posterior distribution, which are concentrated at a relatively fixed position compared to grasps along the rim sampled from the prior.

Hardware implementation. The posterior policy distribution trained in simulation is deployed on the hardware setup (Fig. 5) in a zero-shot manner. We pick 25 mugs with a wide variety of shapes and materials (Fig. A5). Among three sets of experiments with different seeds (for sampling initial poses and latent $z$), the success rates are 100% (25/25), 100% (25/25), and 96% (24/25). The hardware results thus validate the PAC-Bayes bound trained in simulation (Table 1).

In this example, we tackle the challenging task of pushing boxes of a wide range of dimensions to a target region across the table. Real-time external visual feedback is applied using a $150\times 150$ overhead depth image of the whole environment at 5Hz. The observation comprises of the depth map augmented with the proprioceptive x/y positions of the end-effector. The action is the desired relative-to-current x/y displacement of the end-effector $a_{\textup{push}}=(\Delta x,\Delta y)$. A low-level Jacobian-based controller tracks desired pose setpoints. The gripper fingers maintains a fixed height from the table and a fixed orientation, and the gripper width is set to 3 cm to maintain a two-point contact.

Rectangular boxes are generated by sampling the three dimensions (4-8 cm in length, 6-10 cm in width, and 5-8 cm in height) and the mass (0.1-0.2 kg). Each environment again consists of a unique box and a unique initial SE(2) pose. Based on the dimensions of starting and target regions, we define an “Easy” task and a “Hard” one (Fig. 6). A continuous cost is assigned based on how far the COM of the box is from the target region at the end of a rollout.

Expert data. In each of 150 environments, we specify 2 different, successful pushing trajectories. The overhead depth image, the end-effector pose, and the desired end-effector displacement are recorded at 5Hz. It takes about 90 minutes to collect the 300 trials.

Prior performance. Entire trajectories of observations and actions from expert data are embedded into latent $z\in\mathbb{R}^{5}$. Thus $T_{\textup{push}}$ is the total number of steps in each trial. The reconstruction loss $\mathcal{L}_{\textup{rec,push}}$ is again a combination of $l_{2}$ and $l_{1}$ loss between predicted and expert’s actions. The prior policy distribution is able to achieve 84.2% success in novel environments in the “Easy” task, and 74.9% in the “Hard” task. Fig. 4(c) shows a challenging environment of “Hard” task where the box starts far away from the red centerline and with a large yaw angle relative to the gripper. While experts always push at the corner of the box in demonstrations, the prior $P_{0}$ learned via behavioral cloning fails to imitate this behavior perfectly.

Posterior performance. We train both tasks using 500, 1000, and 2000 training environments. The resulting PAC-Bayes bound and empirical success rates across 2000 test environments are shown in Table 2. Posterior performances are improved by about $10\%$ from the priors. Fig. 4(c,d) shows that compared to prior policies, posterior policies perform better at the challenging task as the robot learns to push at the corner consistently.

Hardware implementation. The posterior policy distributions trained using 1000 environments are deployed on the real arm. For both “Easy” and “Hard” tasks, three set of experiments with different seeds are performed with 15 rectangular blocks (Fig. A7). The success rates are shown in Table 2. Note that the result for the “Easy” task falls short of the bound in some seeds. We suspect that the sim2real performance is affected by imperfect depth images from the real camera and minor differences in dynamics between simulation and the actual arm.

In this example, a Fetch mobile robot needs to navigate around furniture of different shapes, sizes, and initial poses, before reaching a target region in a home environment (Fig. 7). We use iGibson [37] to render photorealistic visual feedback in PyBullet simulations. Again we consider a challenging setting where the robot executes actions given front-view camera images and without any extra knowledge of the map. At each step, the policy takes a $200\times 200$ RGB-D image and chooses the motion primitive with the highest probability from the four choices (move forward, move backward, turn left, and turn right), $a_{\textup{nav}}=\operatorname*{arg\,max}(p_{\textup{forward}},p_{\textup{backward}},p_{\textup{left}},p_{\textup{right}})$.

We collect 293 tables and 266 chairs from the ShapeNet dataset [36]. A fixed scene (Sodaville) from iGibson is used for all environments. One table and one chair are randomly spawned between the fixed starting location and the target region in each environment. The SE(2) poses of the furniture are drawn from a uniform distribution. A rollout is successful if the robot reaches the target within 100 steps without colliding with the furniture and the wall.

Expert data. In each of 100 environments, two different, successful robot trajectories are collected. The front-view RGB-D image from the robot and the motion primitive applied at each step are recorded. It takes about 90 minutes to collect the 200 demonstrations.

Prior performance. Short sequences ($T_{\textup{nav}}=3$) of observations and actions from expert data are embedded into latent $z\in\mathbb{R}^{10}$. The reconstruction loss $\mathcal{L}_{\textup{rec,nav}}$ is the cross-entropy loss between predicted action probabilities and expert’s actions. Before each rollout, a single latent $z$ is sampled and then applied as the policy for all steps. The prior policy distribution is able to achieve 65.4% success in novel environments in simulation. Fig. 7(c) shows the diverse trajectories generated by the trained cVAE.

We also investigate the benefit of “fine-tuning” a multi-modal prior distribution vs. a single-modal one. Fig. 7(b,c) show the differences in the policies sampled from these two priors in the same environment. We find that the multi-modality of the prior accelerates PAC-Bayes training and leads to better empirical performance of the posterior. Please refer to A4 for full discussions.

Below are the implementation details for training the posterior distribution with Natural Evolutionary Strategies (NES) (Sec. 3.2). Our objective is to minimize the upper bound $C_{\textup{PAC}}$ in Thm. 1. We will first express the gradient of $C_{\textup{PAC}}$ w.r.t. $\psi:=(\mu,\sigma)$ as follows (using (5) and (6)):

To estimate the expectation on the right-hand side (RHS) of the above equation, we sample a few $z$’s from $\mathcal{N}_{\psi}$ and average $C_{ES}$ over them. In particular, we perform antithetic sampling of $z$ (i.e., for each $z$ that we sample, we also evaluate $C_{ES}$ for $-z$) to reduce the variance of the gradient estimate [38]. This gives us the following ES gradient estimate for $m$ samples of $z$ ($2m$ with antithetic sampling):

Finally, using the Fisher information matrix $F_{\psi}$ (of the normal distribution $\mathcal{N}_{\psi}$ w.r.t. $\psi$) we compute an estimate of the natural gradient [31], denoted by $\tilde{\nabla}_{\psi}$, from the above equation:

The natural gradient estimate computed above is passed to the Adam optimizer [33] to update the belief distribution parameterized by $\psi$. In practice we use $\psi=(\mu,\log\sigma^{2})$ instead of $(\mu,\sigma)$ to avoid imposing the strict positivity constraint on $\sigma$ for the gradient update with Adam.

The derivations follow [27, 28]. Following Sec. 3.2, first the empirical training cost is estimated by sampling a large number of policies $z_{1},...,z_{L}$ from the optimized posterior distribution $\mathcal{N}_{\psi^{*}}$, and averaging over all $N$ training environments in $S$:

Next, the error between $\hat{C}_{S}(\mathcal{N}_{\psi^{*}})$ and $C_{S}(\mathcal{N}_{\psi^{*}})$ can be bounded using a sample convergence bound [34] $\bar{C}_{S}$, which is an application of the relative entropy verision of the Chernoff bound for random variables (i.e., costs) bounded in $[0,1]$ and holds with probability $1-\delta^{\prime}$:

where $\textup{KL}^{-1}$ refers to the KL-inverse function and can be computed using a Relative Entropy Program (REP) [27]. $\textup{KL}^{-1}:[0,1]\times[0,\infty)\rightarrow[0,1]$ is defined as:

The KL-inverse function can also provide the following bound on the true expected cost $C_{\mathcal{D}}(P)$ (Theorem 2 from [28]):

Now combining inequalities (A2) and (A4) using the union bound, the following final bound $C_{\textup{bound}}$ holds with probability at least $1-\delta-\delta^{\prime}$:

As shown in Fig. 7, the prior policy distribution for indoor navigation can exhibit either uni-modal or multi-modal behavior. Uni-modal prior has low entropy while multi-modal prior has higher entropy. We expect that a prior with high entropy can benefit PAC-Bayes “fine-tuning” as opposed to a low-entropy prior; we investigate this further in the indoor navigation example and report the training curves with the low- and high-entropy priors in Fig. A1. While the two priors achieve similar empirical performance before being “fine-tuned”, the prior with higher entropy trains faster and achieves better empirical performance on test environments at the end of “fine-tuning” (0.799 vs. 0.706 in success rate).

As illustrated in Fig. 7, the prior distribution with low entropy always picks the small gap on the left to navigate. Although the policies were successful in behavior cloning environments, they might fail in the “fine-tuning” training environments where the gap can shrink or there can be an occluded piece of furniture behind the gap. Hence “fine-tuning” can be difficult as the prior always chooses that specific route. Instead, the prior with higher entropy “encourages” the robot to try different directions, and is intuitively easier to “adapt” to new environments. “Fine-tuning” then picks the better policy among all available ones for each environment.

All behavioral cloning training is run on a desktop machine with Intel i9-9820X CPU and a Nvidia Titan RTX GPU. PAC-Bayes training for manipulation tasks is performed on an Amazon Web Services (AWS) c5.24xlarge instance that has 96 threads. PAC-Bayes training for the navigation task is done using an AWS g4dn.metal instance that has 64 threads and 8 Nvidia Tesla T4 GPUs. It took about 3 hours to fine-tune the grasping policy with 500 training environments, and about 8 hours for pushing with 500 environments (a GPU instance could have been used to accelerate model inferences). The navigation task is more computationally intensive as it requires GPUs to render the indoor scene - it took about 30 hours with 500 training environments.

All hardware experiments are performed using a Franka Panda arm and a Microsoft Azure Kinect RGB-D camera. Robot Operating System (ROS) Melodic package (on Ubuntu 18.04) is used to integrate robot arm control and perception.