Visually Robust Adversarial Imitation Learning from Videos with Contrastive Learning

We model the decision processes as infinite-horizon discounted Partially Observable Markov Decision Processes (POMDPs) described by the tuple $(\mathcal{S},\mathcal{A},\mathcal{X},\mathcal{T},\mathcal{U},\mathcal{R},\rho_{0},\gamma)$, where $\mathcal{S}$ is the set of states, $\mathcal{A}$ is the set of actions, and $\mathcal{X}$ is the set of observations. $\mathcal{T}:\mathcal{S}\times\mathcal{A}\rightarrow P(\mathcal{S})$ is the transition probability function where $P(\mathcal{S})$ denotes the space of probability distributions over $\mathcal{S}$, $\mathcal{U}:\mathcal{S}\rightarrow P(\mathcal{X})$ is the observation probability function, and $\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ is the reward function which maps state-action pairs to scalar rewards. Finally, $\rho_{0}\in P(\mathcal{S})$ is the initial state distribution and $\gamma\in[0,1)$ the discount factor. The true environment state $s\in\mathcal{S}$ is unobserved by the agent. Given an action $a\in\mathcal{A}$, the next state is sampled such that $s^{\prime}\sim\mathcal{T}(\cdot|s,a)$, an observation is generated as $x^{\prime}\sim\mathcal{U}(\cdot|s^{\prime})$, and a reward $\mathcal{R}(s,a)$ is computed. Note that an MDP is a special case of a POMDP where the underlying state $s$ is directly observed.

Given an MDP and a stationary policy $\pi:\mathcal{S}\to P(\mathcal{A})$, the RL objective is to maximize the expected total discounted return $J(\pi)=\mathbb{E}_{\tau}[\sum_{t=0}^{\infty}\gamma^{t}\mathcal{R}(s_{t},a_{t})]$ where $\tau=(s_{0},a_{0},s_{1},a_{1},\dots)$. A stationary policy $\pi$ induces a normalized discounted state visitation distribution defined as $d_{\pi}(s)=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\mathbb{P}(s_{t}=s|\rho_{0},\pi,\mathcal{T})$, and we define the corresponding normalized discounted state-action visitation distribution as $\rho_{\pi}(s,a)=d_{\pi}(s)\pi(a|s)$. Finally, we denote the state value function of $\pi$ as $V^{\pi}(s)=\mathbb{E}_{\tau}[\sum_{t=0}^{\infty}\gamma^{t}\mathcal{R}(s_{t},a_{t})|S_{0}=s]$ and the state-action value function as $Q^{\pi}(s,a)=\mathbb{E}_{\tau}[\sum_{t=0}^{\infty}\gamma^{t}\mathcal{R}(s_{t},a_{t})|S_{0}=s,A_{0}=a]$. When a function is parameterized with parameters $\bm{\theta}\in\varTheta\subset\mathbb{R}^{k}$ we write $\pi_{\bm{\theta}}$.

Assume we have a set of expert demonstrations $\tau_{E}=(s_{0:T},a_{0:T})$ generated by the expert policy $\pi_{E}$, a set of trajectories $\tau_{\bm{\theta}}$ generated by the policy $\pi_{\bm{\theta}}$, and a discriminator network $D_{\bm{\chi}}:\mathcal{S}\times\mathcal{A}\to[0,1]$ parameterized by $\bm{\chi}$. Generative adversarial IL [8] optimizes the min-max objective

Maximizing (1) with respect to $\bm{\chi}$ is effectively an inverse RL step where a reward function, $r_{\bm{\chi}}(s,a)=-\log(1-D_{\bm{\chi}}(s,a))$, is inferred by leveraging $\tau_{E}$ and $\tau_{\bm{\theta}}$. On the other hand, minimizing (1) with respect to $\bm{\theta}$ can be interpreted as an RL step, where the agent aims to minimize its expected cost. It has been demonstrated that optimizing the min-max objective in (1) is equivalent to minimizing $\mathbb{D}_{\text{JS}}(\rho_{\pi_{\bm{\theta}}}(s,a)||\rho_{\pi_{E}}(s,a))$, so we are recovering the expert state-action visitation distribution [39].

Traditionally, the V-IfO problem assumes that both the expert and the agent operate within the same POMDP. Throughout this paper we relax this assumption and define two different decision processes: namely a target-POMDP for the agent and a source-POMDP for the expert. The target-POMDP is characterized by the tuple $(\mathcal{S},\mathcal{A},\mathcal{X},\mathcal{T},\mathcal{U}_{T},\mathcal{R},\rho_{0},\gamma)$, whereas the source-POMDP is characterized by the tuple $(\mathcal{S},\mathcal{A},\mathcal{X},\mathcal{T},\mathcal{U}_{S},\mathcal{R},\rho_{0},\gamma)$. The primary distinction between these POMDPs lies in their observation probability functions. Despite sharing identical state and action spaces, given the same state $s_{t}$, the expert’s observation $x_{t}^{S}\sim\mathcal{U}_{S}(\cdot|s_{t})$ from the source-POMDP and the agent’s observation $x_{t}^{T}\sim\mathcal{U}_{T}(\cdot|s_{t})$ from the target-POMDP may be different (i.e., we may have $x_{t}^{S}\neq x_{t}^{T}$). We refer to this as visual mismatch.

Given a domain-invariant feature space $\mathcal{Z}$, our AIL pipeline is defined as follows. We initialize two replay buffers $\mathcal{B}_{E}$ and $\mathcal{B}$ to respectively store the sequences of observations generated by the expert and the agent policies, from which we infer the latent state-transitions $(z_{\bm{\delta}},z_{\bm{\delta}}^{\prime})$. We write $(z_{\bm{\delta}},z_{\bm{\delta}}^{\prime})\sim\mathcal{B}$ to streamline the notation. Then, given a discriminator $D_{\bm{\chi}}:\mathcal{Z}\times\mathcal{Z}\to[0,1]$, we write

As mentioned, alternating (2) with an RL step using $r_{\bm{\chi}}(z_{\bm{\delta}},z_{\bm{\delta}}^{\prime})=-\log(1-D_{\bm{\chi}}(z_{\bm{\delta}},z_{\bm{\delta}}^{\prime}))$ leads to the minimization of $\mathbb{D}_{\text{JS}}\big{(}\rho_{\pi_{\bm{\theta}}}(z_{\bm{\delta}},z_{\bm{\delta}}^{\prime})||\rho_{\pi_{E}}(z_{\bm{\delta}},z_{\bm{\delta}}^{\prime})\big{)}$ [40]. Therefore, we are effectively imitating the expert in the latent space $\mathcal{Z}$. Note that the presented AIL can only succeed if $\mathcal{Z}$ is domain-invariant and embeds the relevant goal-completion information necessary to solve the imitation problem. Next, we show how our algorithm C-LAIfO addresses this challenge.

We define the encoder as $\phi_{\bm{\delta}}:\mathcal{X}^{d}\to\mathcal{Z}$, a function mapping sequences of $d\in\mathbb{N}$ observations to the latent space $\mathcal{Z}$. Specifically, we write $z_{\bm{\delta}}=\phi_{\bm{\delta}}(x_{t^{-}:t})$ and $z_{\bm{\delta}}^{\prime}=\phi_{\bm{\delta}}(x_{t^{-}+1:t+1})$ where $t-t^{-}+1=d$. When a data augmentation function $\text{aug}(\cdot)$ is applied to the sequence of observations, we write $\tilde{z}_{\bm{\delta}}=\phi_{\bm{\delta}}(\text{aug}(x_{t^{-}:t}))$ and $\tilde{z}_{\bm{\delta}}^{\prime}=\phi_{\bm{\delta}}(\text{aug}(x_{t^{-}+1:t+1}))$. We train $\phi_{\bm{\delta}}$ to optimize

The steps in (3)–(4) follow the deep $Q$-network optimization pipeline [41, 42], where we add a contrastive auxiliary loss $\mathbb{E}_{\tilde{z}_{\bm{\delta}}\sim\mathcal{B}}[\mathcal{L}(\tilde{z}_{\bm{\delta}})]$ in (3) defined on the encoder $\phi_{\bm{\delta}}$. In (4), the reward function $r_{\bm{\chi}}(z_{\bm{\delta}},z_{\bm{\delta}}^{\prime})$ is computed through the AIL step in (2). Note that we do not perform data augmentation when computing $r_{\bm{\chi}}$ since, during reward inference, we are deploying and not training the encoder $\phi_{\bm{\delta}}$. In practice, adding data augmentation to AIL in (2) decreases the performance.

In (3), the encoder $\phi_{\bm{\delta}}$ is trained together with the critic networks $Q_{\bm{\psi}_{k}}$ ($k=\{1,2\}$) in order to regress $y$ in (4), where $\text{sg}(\cdot)$ stands for stop gradient of the encoder parameters $\bm{\delta}$. The value of $y$ in (4) is computed by summing $r_{\bm{\chi}}$ with the discounted target critic network at time $t+1$. In (4), $a^{\prime}=\pi_{\bm{\theta}}(\tilde{z}_{\bm{\delta}}^{\prime})+\epsilon$ where $\epsilon\sim\text{clip}(\mathcal{N}(0,\sigma^{2}),-c,c)$ is a clipped exploration noise with $c$ the clipping parameter and $\mathcal{N}(0,\sigma^{2})$ a univariate normal distribution with zero mean and $\sigma$ standard deviation. $\bar{\bm{\psi}}_{1}$ and $\bar{\bm{\psi}}_{2}$ are the slow moving weights for the target critic networks. $\mathcal{B}$ is a replay buffer initialized to store interactions $(x^{T}_{t},a_{t},x^{T}_{t+1})$ of the agent with the target environment. Note that the latent state-transitions $(\tilde{z}_{\bm{\delta}},\tilde{z}_{\bm{\delta}}^{\prime})$ are inferred from sequences of observations using $\phi_{\bm{\delta}}$ and, as above, we write $(\tilde{z}_{\bm{\delta}},\tilde{z}_{\bm{\delta}}^{\prime})\sim\mathcal{B}$ to streamline the notation (cf. Appendix B for the complete pseudo-code).

By solving the optimization problem in (3)–(4), our primary goal is to train both the encoder network $\phi_{\bm{\delta}}$ and the critic networks $Q_{\bm{\psi}_{k}}$ to solve the RL problem with reward $r_{\bm{\chi}}$. In other words, this step focuses on retaining the goal-completion information within the latent space $\mathcal{Z}$ such that critic networks $Q_{\bm{\psi}_{k}}$ are successfully learned. We show in our ablation study that backpropagating the gradient from $Q_{\bm{\psi}_{k}}$ to $\phi_{\bm{\delta}}$ is an important step for achieving this goal and solving the imitation problem. Similarly, the types of augmentations performed on the sequences of observations and the choice of auxiliary loss play a crucial role in discarding the visually-distracting information from $\mathcal{Z}$ and dealing with the visual mismatch.

In the following, we introduce the data augmentation techniques and the auxiliary loss $\mathcal{L}$ in (3) for C-LAIfO. We opt for a contrastive method as this leads to good empirical results and good computational efficiency. Contrastive learning constructs low-dimensional representations of high-dimensional data by maximizing agreement between augmented views of the same data example via a contrastive loss in the latent space $\mathcal{Z}$. In our specific case, we define equivalent data as sequences of observations with the same goal-completion information. The data augmentation is randomized by considering a set of pre-determined functions. We will show in our experiments (Section 5) that both the choice of contrastive loss and the set of augmentation functions play an important role in filtering out visually-distracting information.

First, a stochastic data augmentation module transforms any given sequence of observations $x_{t^{-}:t}$ into two views, denoted $\text{aug}(x_{t^{-}:t})_{i}$ and $\text{aug}(x_{t^{-}:t})_{j}$, which are denoted as the positive pairs. Note that two positive pairs must contain the same goal-completion information. Next, the encoder $\phi_{\bm{\delta}}:\mathcal{X}^{d}\to\mathcal{Z}$ extracts representation vectors from augmented sequences of observations. We write $\tilde{z}_{\bm{\delta}}(i)=\phi_{\bm{\delta}}(\text{aug}(x_{t^{-}:t})_{i})$ and $\tilde{z}_{\bm{\delta}}(j)=\phi_{\bm{\delta}}(\text{aug}(x_{t^{-}:t})_{j})$. Finally, we apply the contrastive loss function

where $\mathbbm{1}_{[k\neq i]}\in\{0,1\}$ is an indicator function equal to 1 if $k\neq i$, $\eta$ denotes a temperature parameter, and $\text{sim}(\bm{u},\bm{v})=\bm{u}^{\intercal}\bm{v}/||\bm{u}||||\bm{v}||$ denotes the cosine similarity. We sample a batch of $N$ sequences of observations from the buffer $\mathcal{B}$ and define pairs of augmented sequences derived from the batch, resulting in $2N$ data points. The negative data points are not sampled explicitly. Instead, given a positive pair, we treat the other $2(N-1)$ augmented data points within the batch as negatives. Note that the loss in (5) is called the normalized temperature-scaled cross entropy loss or the Information Noise-Contrastive Estimation (InfoNCE) loss [43] and represents an upper bound of the negative mutual information between positive pairs. Therefore, by minimizing (5) as in (3) we are maximizing the mutual information between positive pairs in the latent space $\mathcal{Z}$.

Given the environments in Fig. 2, we perform the following ablations:

1. 1.

   Contrastive loss function: We demonstrate the importance of the contrastive loss function type by comparing the InfoNCE loss in (5) with BYOL in [4].

2. 2.

   Gradient backpropagation: We highlight the necessity of backpropagating the gradient from $Q_{\bm{\psi}_{k}}$ to $\phi_{\bm{\delta}}$ in (3) for solving the imitation problem and embedding the goal-completion information in the latent space $\mathcal{Z}$.

3. 3.

   Data augmentation: We emphasize the importance of selecting the appropriate augmentation for a given mismatch, showing that a mismatch-informed augmentation outperforms general augmentations or no augmentation.

The results of these experiments are summarized in Table 1, which includes results for C-LAIfO and its various configurations. All the learning curves are provided in Appendix D. In Table 1, C-LAIfO is implemented as described in Section 4, with the data augmentation function $\text{aug}(\cdot)$ defined as a brightness transformation. In BYOL-LAIfO, we retain the design identical to C-LAIfO except for replacing the InfoNCE loss in (5) with BYOL [4]. In C-LAIfO w/o $Q$ backprop, we disable gradient backpropagation from $Q_{\bm{\psi}k}$ to $\phi_{\bm{\delta}}$ in (3). Lastly, in C-LAIfO full aug and C-LAIfO w/o aug, we respectively modify the data augmentation function $\text{aug}(\cdot)$ to include full augmentation (brightness, color, and geometric transformations, as detailed in Appendix 5) and no augmentation. Notably, without backpropagating the gradient from $Q_{\bm{\psi}_{k}}$ to $\phi_{\bm{\delta}}$, C-LAIfO fails to solve the imitation task even in the simplest mismatch scenario in (2(b)). This result demonstrates the importance of this step for embedding goal-completion information in $\mathcal{Z}$. Conversely, without a well-designed $\text{aug}(\cdot)$, C-LAIfO can solve the imitation problem with minor visual mismatch but struggles with more challenging visual distractors. These results highlight the critical role of properly defining $\text{aug}(\cdot)$ for efficient visual generalization in $\mathcal{Z}$. General $\text{aug}(\cdot)$ might still achieve this goal, but much less efficiently. Finally, the superior performance of the InfoNCE loss in (5) compared to BYOL is particularly evident in handling the most difficult mismatches in (2(e)) and (2(f)).

In this section, we test C-LAIfO in the V-IfO setting with different types of mismatches (cf. Fig. 3) and compare it with three baselines: LAIfO [1] and PatchAIL [2], both equipped with the same $\text{aug}(\cdot)$ used for C-LAIfO, and DisentanGAIL [6]. The results, summarized in Table 2, demonstrate that C-LAIfO successfully addresses the V-IfO with mismatch problem, achieving superior performance compared to all the baselines across the proposed mismatches. All the learning curves are provided in Appendix D. For the Light experiment, $\text{aug}(\cdot)$ is defined as a brightness transformation; while for the others it is defined as a color transformation. Details on $\text{aug}(\cdot)$ are provided in Appendix 5. Furthermore, to assess whether C-LAIfO achieves domain invariance in the feature space $\mathcal{Z}$, we perform PCA and t-SNE analyses on the latent space $\mathcal{Z}$ learned by different algorithms during training. Specifically, focusing on the Light experiment in Table 2, Fig. 4(a) compares the latent space $\mathcal{Z}$ learned by LAIfO in [1], C-LAIfO, and LAIfO with data augmentation. Additionally, Fig. 4(b) investigates whether the latent space learned by C-LAIfO during the Full experiment can generalize to the unseen environment depicted in Fig. 3(g). The results of this analysis are shown in Fig. 4, which includes the PCA visualization. For the t-SNE visualization and additional details refer to Appendix E.

In the following section, we examine a scenario where the reward $r_{\bm{\chi}}$ learned by C-LAIfO can be combined with a sparse reward $\mathcal{R}$ that the agent collects through interaction with the environment. Consequently, the RL problem aims to maximize the total reward $\mathcal{R}_{\text{tot}}=\mathcal{R}(s_{t},a_{t})+r_{\bm{\chi}}(z_{t},z_{t+1})$ with $r_{\bm{\chi}}$ learned through the AIL step in (2). We conduct experiments on the Adroit platform for dynamic dexterous manipulation [7] and compare C-LAIfO with DrQ-v2 [45], an RL-only baseline. In these experiments, both C-LAIfO and DrQ-v2 use an encoder to process pixel observations and extract embeddings in $\mathcal{Z}$, which are then concatenated with the robot sensory observations. Importantly, expert’s sensory observations are not used in the imitation problem as we only assume access to videos of experts. We denote our approach as RL+C-LAIfO since it maximizes $\mathcal{R}_{\text{tot}}$ rather than only $r_{\bm{\chi}}$ alone as in the standard imitation learning problem. The results, summarized in Table 3, show that DrQ-v2 fails to learn in two out of three environments. In contrast, C-LAIfO successfully leverages expert videos with visual mismatch to facilitate learning. Note that the goal of these experiments is not merely to compare RL+C-LAIfO with DrQ-v2 but to demonstrate how C-LAIfO can leverage expert videos, even with visual mismatch, to enable learning in challenging robotic tasks. We provide all the learning curves in Appendix D.

A main limitation of the current approach is given by the reliance of C-LAIfO on a well-designed, possibly mismatch-informed, data augmentation function. As illustrated in our ablations in Section 5, general augmentation can lead to poor performance or can remarkably reduce the algorithmic sample efficiency. Furthermore, it can be challenging to design effective augmentations for certain types of mismatches. To address this problem, exploring generative models for automatic data augmentation represents an interesting research direction. Generative models could produce diverse, mismatch-informed augmentations, potentially overcoming the limitations of manually designed strategies. Alternatively, investigating different auxiliary losses that are less reliant on augmentation techniques represents another interesting direction. Finally, future work will be devoted to go beyond simulated environments and test our algorithms on hardware in real-world scenarios.