Towards More Generalizable One-shot Visual Imitation Learning

Figure 2 provides an overview of our model pipeline: a CNN backbone is followed by a multi-head self-attention module [41] and an MLP to make action predictions.

Visual features Given a batch of $B$ inputs containing $T_{d}$ expert video frames and $T_{o}$ agent observations, a CNN backbone encodes each frame into $C$ channels of size $H\times W$ feature maps, resulting in demonstration features $\mathbf{x}_{\mathit{d}}$ with size $[B,C,T_{d},H,W]$, and observation features $\mathbf{x}_{o}$ with size $[B,C,T_{o},H,W]$. To preserve spatial and temporal information, both features are flattened along the last 3 dimensions and added with sinusoidal encodings [41], then re-shaped into the original size.

Self-attention module We use multiple self-attention layers that model the underlying relationship between the sequence of representations $\mathbf{x}_{\mathit{d}}$ and $\mathbf{x}_{o}$. We adopt the non-local self-attention block in [42], and make it to a multi-head version as [7]. Specifically, key, query and value tensors are generated from three separate 3D convolution layers, which are then flattened along the time and space dimensions to compute spatio-temporal attention by each head individually. Formally, given temperature parameter $\tau$, key $K_{j}$, query $Q_{j}$ and value $V_{j}$, the attention head $j$, the attention operation is computed as:

$\displaystyle A_{j}=\mathbf{softmax}(Q_{j}K_{j}^{\top}/\tau)V_{j}.$

Outputs $A_{j}$ from each head are concatenated and projected to the original feature size by another 3D convolution, as used in [42]. $\mathbf{x}_{\mathit{d}}$ will first pass through a self-attention module to get $\mathbf{x}_{\mathit{d}}^{\text{attn}}$. Then, every frame in $\mathbf{x}_{o}$ will compute self-attention with both $\mathbf{x}_{\mathit{d}}^{\text{attn}}$ and itself, which in effect calculates: (1) the spatial self-attention in each observation frame and (2) spatio-temporal cross-attention on the demonstration. The resulting $\mathbf{x}_{o}^{\text{attn}}$ will be used to predict action.

Our method bases off the intuition that, representations for two nearby frames from the same video clip should be similar, whereas frames from different tasks or variations should be drawn apart. For each frame in a video, we maximize its feature similarity with a randomly selected, temporally close-by frame. Specifically, we take an input batch, and obtain its two “view”s by two separate data augmentations. The model encodes the first view into $\mathbf{x}_{1}$, and a target model encodes the second to get $\mathbf{x}_{2}$, which is gradient-free and receives parameter updates solely from its online counterpart. Lastly, $\mathbf{x}_{1}$ and $\mathbf{x}_{2}$ are separately passed through a linear projector $f$: $z_{1}=f(\mathbf{x}_{1})$, and its target $\bar{f}$: $z_{2}={\bar{f}}(\mathbf{x}_{2})$. For every feature frame in $z_{1}$, we select a nearby feature frame in $z_{2}$ as positive. We then maximize the similarity between each anchor $q=g(z_{1})$, and its positive $k_{+}=z_{2}$, via the InfoNCE loss from [40], where $g$ is another linear projector, also named as predictor in prior work [12].

We follow [40, 15] to model embedding similarity as bilinear product, calculated with a projection matrix $W$. Formally, with total frame count $F=B(T_{d}+T_{o})$, treating every other $k$ in the batch as negatives, the contrastive loss at each $q$ is expressed as:

$\displaystyle\mathcal{L}_{\text{Rep}}=\log\frac{\exp\left(q^{T}Wk_{+}\right)}{\exp\left(q^{T}Wk_{+}\right)+\sum_{i=1}^{F-1}\left(q^{T}Wk_{i}\right)}$

(1)

One may view the convolution backbone and self-attention layers as one combined feature extractor, therefore the above contrastive loss can be applied to either before or after the self-attention layers as shown in Figure 2. Moreover, our contrastive module differs from [23] in the new temporal contrast strategy: for frame feature $x_{t}$ at timestep $t$, it contrasts with a random nearby frame selected from $x_{t-k}$ to $x_{t+k}$ in its augmented counterpart, whereas prior work [37] uses a fixed-step future frame. We provide ablation experiments in Appendix to provide additional insights on details of our contrastive objetive implementation.

Our objective is to learn a policy $\pi_{\theta}(a_{t}|o_{t},\mathit{d})$ which takes current image observation and a demonstration video as inputs, and predicts the action distribution to successfully finish the task.

To enable learning a potentially multi-modal policy that excels across many tasks, we adopt the same solution used in [7, 25, 33], which discretizes the action space into 256 independent bins along every dimension, and parameterize the policy using a mixture of discretized logistic distribution. As described in Section IV-A, the self-attended observation features $\mathbf{x}_{o}^{\text{attn}}$ will pass through the action MLP, to predict the mean $\mu_{i}$, scale $s_{i}$ and mixing weight $\alpha_{i}$ for each discretized logistic distribution. The behavior cloning training loss is the negative log-likelihood:

$\displaystyle\mathcal{L}_{\text{BC}}=-\log\left(\sum_{i=1}^{m}\alpha_{i}(\mathbf{x}_{o}^{\text{attn}})P\left(a_{t},\mu_{i}(\mathbf{x}_{o}^{\text{attn}}),s_{i}(\mathbf{x}_{o}^{\text{attn}})\right)\right)$

(2)

Where $P\big{(}a_{t},\mu_{i},s_{i}\big{)}=\sigma(\frac{a_{t}+0.5-\mu_{i}}{s_{i}})-\sigma(\frac{a_{t}-0.5-\mu_{i}}{s_{i}})$, $\sigma$ is the logistic sigmoid function. At the inference time, given $o_{t}$, the action is sampled from the predicted distribution:

$\displaystyle a_{t}\sim\sum_{i=1}^{m}\alpha_{i}(\mathbf{x}_{o_{t}}^{\text{attn}})\text{ logistic}\left(\mu_{i}(\mathbf{x}_{o_{t}}^{\text{attn}}),s_{i}(\mathbf{x}_{o_{t}}^{\text{attn}})\right)$

(3)

In addition to the behavioural cloning loss, we also utilize the inverse dynamics loss as in  [7]. By taking consecutive observation frames $o_{t},o_{t+1},\dots,o_{t+k}$ during training, another MLP will predict inverse actions $a_{t},a_{t+1},\dots,a_{t+k-1}$. The inverse dynamics loss has similar form as (2):

	$\displaystyle P_{i}=P\left(a_{t},\mu_{i}(\mathbf{x}_{o_{t}}^{\text{attn}},\mathbf{x}_{o_{t+1}}^{\text{attn}}),s_{i}\left(\mathbf{x}_{o_{t}}^{\text{attn}},\mathbf{x}_{o_{t+1}}^{\text{attn}}\right)\right)$		(4)
	$\displaystyle\mathcal{L}_{\text{Inv}}=-\log\left(\sum_{i=1}^{m}\alpha_{i}(\mathbf{x}_{o_{t}}^{\text{attn}},\mathbf{x}_{o_{t+1}}^{\text{attn}})P_{i}\right)$		(5)

Combing with the contrastive loss $\mathcal{L}_{\text{Rep}}$ introduced in Section IV-B, we obtain the overall loss for our method:

$\displaystyle\mathcal{L}=\lambda_{\text{Rep}}\mathcal{L}_{\text{Rep}}+\lambda_{\text{BC}}\mathcal{L}_{\text{BC}}+\lambda_{\text{Inv}}\mathcal{L}_{\text{Inv}}$

(6)

Simulation environment. We develop 7 distinct task environments using Robosuite v1.1 [48] and combining MetaWorld [46] for additional assets. For each task, we additionally design multiple semantically distinct variations. In order to investigate cross-morphology imitation, we also integrate two robot arms. The imitation policy is learned and evaluated on a Panda robot arm but takes a Sawyer robot video as demonstration.

Data collection. For every variation of each task environment, we design scripted expert policies and collect 100 demonstration videos of Sawyer robot and another 100 for Panda robot, with differently initialized scene layouts as instances. We provide more detailed information on simulation environment and data collection in Appendix.

We conduct experiments with the dataset described in Section V-A to answer the following questions:

• •

  How does our method compare with prior baselines under the original single-task one-shot imitation learning setup.

• •

  How well does our method perform across multiple tasks, after trained on the same set of tasks.

• •

  How well does our trained multi-task model perform given a completely new task: can it 1) directly perform one-shot imitation at test time; 2) be fine-tuned to quickly adapt to the new task requiring fewer amount of data.

• •

  Which component(s) in our contrastive module are key to its effectiveness at representation learning. The ablation experiment results are provided in Appendix due to space limitation.

We report and compare performance to the following baseline methods:

• •

  DAML [45]: We train a policy model with MAML [10] loss and behaviour cloning loss. We replace the model architecture used in the original paper with a wider and deeper network of comparable parameter counts as ours, and use the same action distribution parametrization for policy learning.

• •

  T-OSIL [7]: Model architecture proposed by [7] which also uses non-local block [42] in self-attention module and is trained with end-effector point prediction as auxiliary loss. Our method utilizes a more computationally efficient attention operation, see Figure 12 for an illustrated comparison.

• •

  LSTM: We replace the self-attention module in our model architecture with linear projectors followed by an LSTM [17, 38] architecture of a comparable parameter count. The rest of the policy model architecture is kept the same as in our main method.

• •

  MLP: We replace the self-attention module in our model architecture with a simple MLP layer to process stacked visual features from the demonstration into “task context vectors”, which is then concatenated with observation features and used for action predictions.

For all experiments (including baselines), we keep the first three convolutional residual blocks in ResNet-18 as the feature extractor, and apply the same data augmentation strategy to prevent over-fitting and improve model robustness. For evaluation, we take 3 different converged model checkpoints, and for each variation of each task, we gather each policy’s rollout performance across 10 episodes with different random seeds. For both the single and multi-task models, we report the mean and standard deviation of success rates in each task separately. Details on network architecture and hyper-parameters can be found in the Appendix.

Single-task One-shot Imitation

We first evaluate performance on the single-task setup as done in prior methods. Specifically, the model is trained with demonstrations from multiple variations of one task and then tested on unseen instances of the same task (e.g. different initial poses). We report the success rate of all methods on each task in the rows named “single” of Table I, with a clear out-performance of our method over baselines on every task.

We remark that DAML [45] also experimented with a pick-and-place task in their original paper, but collected “hundreds of” objects for training and 12 held-out objects for testing. Without access to further details, we hypothesize that this visual diversity in training dataset was crucial for its success at picking correct test-time objects, which explains why the same method reports massive under-performance on the 4-object Pick & Place task, where [7] also reported a very low success rate from DAML ($6.9\%$) using a differently-configured simulation of the same task.

Multi-task One-shot Imitation

We next consider the multi-task setup by mixing data from all 7 tasks into one training dataset. After training, we report the success rate of one model on each task in the rows named “multi” of Table I. The baselines’ performances drop significantly compared with their results from single-task, whereas ours continues to work well across many tasks ¹¹1For this setup, we increase the number of attention layers in the model from 2 to 3, and also adjust each baseline accordingly for fair comparisons. However, the performance of DAML [46] is still not comparable with others using the updated architecture and after tuning hyper-parameters..

Novel Task Generalization

To show a higher level of generalization ability, we test a one-shot imitation agent with tasks that are sufficiently different from what it already trained on. We hence set up a series of experiments that, each picks 1 out of the 7 tasks in our benchmark suite as the held-out task, and trains a multi-task model on the remaining 6 tasks until convergence.

We first directly evaluate each model on its corresponding held-out task. Results are reported in column “Novel-task” of Table II, where each row corresponds to the experiment where the current task was excluded from training and only used for one-shot evaluation. For comparison purposes, we also include: “No Training” column, which evaluates a randomly initialized policy network without any training , “Single-task” column, where each row reports performance of a model trained and tested on the same task, and “Multi-task” column, where we take one model trained on all 7 tasks and report its performance on each task separately.

As shown in Table II, directly evaluating a multi-task model largely fails to complete an unseen novel task, and performs significantly worse than when this novel task was included during single- or multi-task training. This failure of direct one-shot imitation on a novel task suggests exciting room for future research. Nevertheless, the gap between “Novel-task” and “No Training” results ²²2 We remark that, the non-zero success rates in “No Training” column are due to the nature of task design in simulation, where a randomly initialized policy model would sometimes generate an action that accidentally leads to an episode being counted “successful”, such as hitting an opened drawer to shut it close or stumbling on a closed door and pushing it open. suggests a multi-task model still learns certain non-random behaviors, which could be generalized directly to a completely novel task. Note that, the improvement over “No Training” is limited to the first three tasks (Door, Drawer and Press Button), which require simpler motions and task reasoning.

We are hence encouraged to continue from this setup, but further fine-tune each 6-task model on its corresponding held-out task. We use the exact same dataset for training the models in “Single-task” column, and report the final fine-tuned results in “Fine-tune” column. To investigate the data efficiency of this setup, we additionally experiment with using $25\%$, $50\%$, $75\%$ the amount of the original single-task training data for fine-tuning, and compare with training from scratch on that single task using the same dataset size.

We plot evaluation results of intermittently-saved checkpoints during each model’s training in Figure 3. We observe that a multi-task pre-trained model is able to adapt quickly to a completely new task, even when limited data is available (i.e., $25\%$), and the final convergence performance is sometimes higher than training single-task from-scratch for some challenging tasks (e.g. Nut Assembly, Block stacking and Basketball). This indicates that these models are indeed able to accumulate generalizable knowledge (such as flexible visual feature extractors) from pre-training on many tasks, which then put them at advantage of learning a novel task very efficiently.

In this section, we first provide detailed descriptions of our task environment design, then introduce more details about our method implementation and the experiment setup, and lastly we provide more ablation experiment results and analysis.

We use Robosuite [48] as our base framework, and design 7 robot manipulation environments (tasks) with intra-task variations. For example, consider Nut Assembly task in the bottom-right of Figure 11: it has 9 variations in total, resulting from picking up any one of the 3 nuts on the table, then assembling it to one of the 3 pegs. Since there are multiple varied instances for a specific task, the agent should finish the correct task without mis-identification based on demonstration. Below we provide details about each of the individual task design.

Network Architecture We provide an illustration of our self-attention block and one of our baseline [7] as in Figure 12. For other baselines in our main paper, we replace this attention module with LSTM or MLP accordingly.

Data Augmentation We add data augmentation for all the baselines to prevent overfitting. In our implementation, we use 4 types of data augmentation provided in the torchvision package: random translate, random crop, color jitter, Gaussian blur.

Hyper-parameter Settings We provide more details about the hyper-parameters and other settings of model training and evaluation in Table III.

Batch Construction One batch consists of demonstrations sampled from each variant in every training task, which are mixed evenly. Since each task contains a different number of variations, to ensure they have comparable learning progress, the loss is first averaged across variants within each task, and then across different tasks.

Computation Requiremenmts Our single-task model can be trained within one GPU day of NVIDIA TITAN Xp, whereas the multi-task model takes about 4 GPU days to train.

We conduct a set of ablation studies focused on our contrastive representation learning approach, to provide insights on the design of the loss objective, how it incorporates with the policy learning, and its impact over the one-shot imitation performance.

The effect of contrastive learning objective

We first ablate by removing contrastive loss from the model’s training update: we train a model with solely behavior cloning loss, while other training configurations are kept consistent with the multi-task experiments reported in Table I. Results are reported in row “Single-task No Contra.” and “Multi-task No Contra.” of Table IV. Noticeably, our multi-task performance is significantly decreased without the contrastive learning objective. Comparing results from single-task and multi-task setups, we observe a clear challenge of task reasoning and robust representation learning, which the contrastive objective is able to address, but not fully closing the gap between a multi-task model and its single-task counterpart on every task.

Adding the contrastive loss also shows a more significant gain on the right-most 3 tasks of Table IV, which vary only the colors of otherwise similarly-shaped objects. As compared to Pick & Place where the 4 objects are differently sized, shaped, and colored, the blocks/basketballs/nuts in these variations tend to be more easily confused, which can be addressed the contrastive loss that explicitly forces the representations to be distinguishable among different sub-tasks.

To investigate the effect of contrasting against negative samples in the loss objective design, we implement BYOL [12], in which a data sample is only drawn together with its augmented counterpart. Its multi-task performance is shown in Table IV, which show no improvement over not using contrastive loss at all (as reported in row “Multi-task No Contra.”). We remark that in multi-task setups, negative samples help learn stronger representations that are better at distinguishing among images from different tasks/variations.

Where to apply contrastive loss. As discussed in Section IV-B, we can apply the contrastive operation in features from any intermediate layers of the model. To understand the effects of this algorithmic choice, we consider the following variants of our method: (1) Pre-Attn: applying contrastive loss to only features prior to the self-attention layers; (2) Post-Attn: applying contrastive loss only after the self-attention layers; (3) Both-Attn (Ours): calculating both losses in (1) and (2), which we use for single and multi-task experiment results in above sections. Table V shows the performance of each ablation variant in three single tasks, among which we find that Both-Attn (using two contrastive losses) achieves the best performance.

The effect of temporal contrast.

Our contrastive strategy is doing a temporal contrast by randomly selecting frames from nearby time-steps as positive. To fully understand its effectiveness, we compare it with two variants: (1) No-Temp: applying contrastive loss to two different data augmentations from same frame, similar to CURL [23]; (2) Fix-Temp: applying temporal contrast but always with a fixed-step future frame, similar to ATC [dwibedi2019temporal]. Ours (denoted Rand-Temp) achieves the best performance as shown in Table V, as suggested by comparing the first, second and last column in Table V.

In this section, we provide a more detailed discussion of the methods and experiment task settings used by related prior work in one-shot imitation learning (OSIL).

However, prior OSIL work has been limited to a single-task setup and mainly tests a model on a different task variation (e.g stacking an unseen block combination) or a different instance (e.g. different object pose) of the previously-seen variations. Experiments in [9] train an agent to stack various (unseen) block combinations at test time, but use low-dimensional state-based inputs. For visual inputs, [11] and [21] experimented with 3 separate settings: simulated planer reaching (with different target object colors), simulated planer pushing (with varying target object locations), and real-robot, object-in-hand placing (onto different target containers); [45] set up a two-stage pick-then-place task with varying target objects and target containers; [7] uses a simulated Pick & Place task with 4 objects to pick and 4 target bins to place (hence 16 variations in total). The AI2-THOR [22] environment used in [19] requires collecting varying objects and dropping off at their designated receptacles, where actions are purely semantic concepts such as “dropoff” or “search”. In contrast, in this work we consider a harder, multi-task setup, where agent needs to perform well across more diverse and distinct tasks, and generalize not only to new instances of all the seen variations, but also to completely novel tasks.