Unsupervised Feature Learning for Manipulation with Contrastive Domain Randomization

Domain randomization is a class of domain adaptation techniques [12, 13] that targets the transition of learned models from simulation to the real-world. DR has been successful on vision problems (e.g. object detection, pose estimation [14]) often using a rendering engine to perturb non-essential properties in synthetic scenes [15]. Recent work has extended DR to robotic grasping [16], navigation [17], locomotion [18], and, relevantly, deformable object manipulation [19]. Further methods alter simulation dynamics in addition to visual features [20, 21], randomizing physical properties, such as friction and latency, to learn robust manipulation policies. Implicit in DR is an assumption that the environment is composed of relevant and irrelevant properties, which together generate an observation. The goal is to learn a representation that attends to relevant properties (e.g. physics) while ignoring the irrelevant ones. In this work, we make a connection between the domain randomization objective and the idea of learning invariances, as studied in the causal inference literature [22]. Recent work on invariant risk minimization [23] also aims to learn invariant features. However, our setting controls the simulated domain (the ‘intervention’ in causal inference parlance), enabling more effective representation learning.

Contrastive learning has gained popularity as a method for self-supervised feature learning, and relates to information-theoretic concepts such as mutual information [5, 6, 7]. Leveraging spatially- or temporally-coherent data (e.g., images or videos), contrastive learning techniques construct ”pretext” tasks [24] in which a neural network must discriminate between similar and dissimilar input samples [25, 26]. Pretext tasks can involve predicting image patches [27, 5], frames in videos [5], optical flow [6], programs [28], and other visual features [29, 26, 25]. Crucially, contrastive learning can be interpreted as learning a non-linear transformation that is invariant to distortions of the input data [30]. Our method combines these ideas with domain-randomization to learn self-supervised domain-invariant features.

We focus on deformable object manipulation approaches that combine predictive models with forward planning. Recent work has proposed learning non-linear deformation functions to assist feedback control [31, 32]. Generative modeling [33, 34] has enabled neural networks to learn richer dynamics that can be used by model-predictive control (MPC) for manipulation [35, 36]. Contrastive learning methods, which represent spatial/temporal patterns in the data as latent vectors, are particularly useful for model-based planning [37, 38]. For example, recent methods applied MPC for cloth and rope manipulation, using a forward model learned using contrastive learning [39, 11]. Our work extends these approaches to planning that is invariant to irrelevant visual properties.

Contrastive Learning methods learn compact representations of high dimensional observations such as images, video, or audio. Unlike supervised learning, where a parametric model $f$ is trained to maximize the similarity between a prediction $\hat{z}$ and a known label $y$, contrastive learning trains an encoder $f_{\theta}$, parametrized by $\theta$, to map observations $o$ into a latent representation $\hat{z}=f_{\theta}(o)$, and makes predictions directly in the latent space. For example, Contrastive Predictive Coding (CPC) [5] introduces a self-supervised instance discrimination task in which $f_{\theta}$ is used to predict a single positive label $y$ amongst $N-1$ contrastive labels $y_{j}$. This objective, known as InfoNCE loss, simultaneously maximizes the similarity $h(\hat{z},y)$ between an observation based prediction $\hat{z}$ and the positive label $y$ while minimizing the similarity $h(\hat{z},y_{j})$ between $\hat{z}$ and contrastive labels $y_{j}$. As a self-supervised task, positive labels $y$ are computed from the observation itself, while contrastive labels $y_{j}$ are computed from different observations in the data. InfoNCE frames the discrimination task as a standard classification problem, through the cross-entropy loss:

Typically, the InfoNCE loss is minimized using stochastic gradient descent, where a batch contains random observations from the data, and the loss for a batch of size $N$ is:

That is, we set the constrastive labels for an observation to be the labels of other observations in the batch.

Contrastive samples often leverage spatial, temporal, or semantic consistency in the input data. Given image data $o$, for example, it is possible to generate a prediction $\hat{z}=f_{\theta}(\tau^{\prime}(o))$ and positive label $y=f_{\theta}(\tau(o))$ by applying separate visual transformations ($\tau^{\prime}$ and $\tau$, respectively) to a single image $o$. Contrastive labels can then be generated by applying transformations to different images $y_{j}=f_{\theta}(\tau(o_{j}))$ [7, 25]. For sequential data, the temporal dimension of the input serves as a useful signal with which to generate contrastive labels. CPC follows this approach, defining the instance discrimination task over frames in a video [5]. Specifically, from the latent representation $z_{t}$ of the image at time $t$, $o_{t}$, CPC predicts the latent representation of a frame $k$ steps in the future $\hat{z}_{t+k}$, using the representation of the true frame $z_{t+k}$ as a positive label, and representations of frames from random time-steps (or different videos altogether) as contrastive labels $z_{j\neq t+k}$. In doing so, CPC simultaneously maximizes the similarity $h(\hat{z}_{t+k},z_{t+k})$ between the predicted and true future frame representations and minimizes the similarity $h(\hat{z}_{t+k},z_{j\neq t+k})$ between the prediction and each of the contrastive sample representations. Key to this approach is the joint optimization of an encoder $z_{t}=f_{\theta}(o_{t})$ and an auto-regressive model $\hat{z}_{t+k}=g_{\phi}(z_{\leq t})$, parametrized by $\phi$, that at time $t$, produces a latent prediction $\hat{z}_{t+k}$, given a summary $z_{\leq t}$ of observations up to time $t$. Our method similarly exploits the temporal structure of video data for contrastive sampling.

Several similarity metrics have been proposed in the literature. A common choice is the weighted dot-product $\exp(z_{i}^{T}Wz_{j})$ or $\sigma(z_{i}^{T}Wz_{j})$ [7, 25]. The cosine distance and L2 distance have also proven effective [5, 11, 40].

Domain randomization is a popular method of simulation-to-real transfer for neural networks [17, 14]. By exposing the network to extensive variations of scene parameters that are irrelevant for decision making during training (e.g. lighting conditions and textures), the goal of DR is to learn robust features that will transfer well to the real world. Here, we cast DR under a probabilistic formulation that will serve our subsequent development.

We assume that an observation (either simulated or real) $o(x,e)$ is generated from two independent random variables $x$ and $e$, representing relevant and irrelevant domain properties, respectively. We further assume access to the observation generation process—i.e. given $x$ and $e$, we can generate $o(x,e)$. For example, $x$ can represent the pose of an object, $e$ its texture, and $o(x,e)$ is generated by a rendering engine. Under these assumptions, DR can be framed as a supervised learning problem over the following loss: $\mathcal{L_{\text{DR}}}(l)=\mathbb{E}_{e}\mathbb{E}_{x}\mathbb{E}_{y|x,e}\left[l(o(x,e),y)\right],$ where $y$ is a label for the observation $o(x,e)$, sampled from $p(y|x,e)$, and $l$ is some supervised learning loss function, such as regression or classification. If $y$ is independent of $e$, optimizing over this loss encourages the network to learn a prediction that is agnostic to $e$.

Recall that observations $o(x,e)$ are generated from two independent variables—relevant properties $x$ and irrelevant environmental properties $e$—and $z=f_{\theta}(o(x,e))$ is an encoded representation of $o(x,e)$. As in DR, we assume access to the observation generation process. Let $y(x,e)$ be a random variable we wish to predict from $o(x,e)$, such as the outcome of an action $a_{t}$ or an encoding of future observations. This framework supports two types of causal models: (i) domain-dependent models in which $y$ is influenced by both $e$ and $x$ (Figure 1a); (ii) domain-independent models in which $x$ predicts $y$, regardless of interventions over $e$ (Figure 1b). In the next section, we show that straightforwardly applying domain randomization to the CPC objective results in domain-dependent representations.

A simple application of DR to CPC is to domain-randomize each sample in our data, and apply CPC on this data [11]. The corresponding loss for a batch of size $N$ can be written as:

where $x_{i},e_{i}$ correspond to random observations $o(x_{i},e_{i})$ in the data, and we denote by $\hat{z}(x_{i},e_{i})$ and $y(x_{i},e_{i})$ the prediction and label that correspond to the observation.

Though Eq. (1) suggests that CPC maximizes the mutual information between the representation $z$ and the observation $o$, when inspecting Eq. (2), we note that there is nothing in the loss function preventing both $y$ and $z$ from becoming dependent on $e$. That is, Eq. (2) may learn to maximize mutual information between $e$ and $z$, effectively learning features that encode spurious properties of the random domain. In fact, depending on the task, $e$ might serve well in distinguishing positive labels from contrastive labels! Indeed, we empirically show that this naive CPC formulation learns representations that depend on the random domain.

An alternative approach to Eq. (2) is to contrast between different observations in the data that share the same domain $e$. This will prevent CPC from learning to discriminate between labels based on $e$. The corresponding loss for a batch is:

where $x_{i},e$ correspond to random observations from the data that share the same domain. While this approach is appealing, it suffers from several technical problems. The first is that SGD based optimization tends to work well when samples in a batch are not correlated. The second is that this method requires generating at least $N$ different observations for each combination of domain properties $e$, which becomes prohibitively large when we want high variation in $e$.

We propose a novel contrastive loss function that learns domain-invariant representations while preserving high sample efficiency. Our method exploits domain randomization to remove the dependence between predictive variables $y$ and spurious domain features $e$.

We now show that CDR indeed learns domain invariant representations. We focus on neural-network encoders, and specifically assume that the encoder can be written as: $f(x,e)=\sigma(W_{x}^{T}\phi_{x}(x)+W_{e}^{T}\phi_{e}(e)+b),$ that is, the encoder has separable features for $x$ and $e$. In this case, the following proposition shows that the CDR loss will learn to ignore $e$.

If $f$ learns to produce $\phi_{x}(o(x,e)){=}\phi_{x}(x)$ and $\phi_{e}(o(x,e))=\phi_{e}(e)$ in any layer of the neural network, then according to Proposition 1 the optimal representation will be domain-invariant. While the assumption on separable features is idealistic, modern neural-network architectures are often expressive enough to be capable of separating $x$ and $e$. In such cases, the result above guarantees invariant features. In our experiments, we further demonstrate this result empirically.

In the uncontrolled setting we only observe sequences of images $\{o_{t}\}$ depicting some physical interaction.

We follow CPC [5] and define $g_{\phi}$ to be an auto-regressive model (GRU) that summarizes the past observations in the latent space $z_{\leq t}$ and predicts the next $K$ future latent states $\hat{z}_{i}=\hat{z}_{t+k}=W_{k}^{T}GRU(z\leq t)$. we set our positive labels $y^{\prime}$ to be the true representation at time $t+k$, $z^{\prime}_{i}=z^{\prime}_{t+k}=f_{\theta}(o(x_{t+k},e^{\prime}))$, and contrastive labels are drawn from different time-steps or from different videos. we average $K$ predictions and optimize:

In the controlled settings, an agent interacts with the environment and we can observe its actions. That is, we observe $o_{t}$ and an action $a_{t}$ at each time-step. Here, we define $g_{\phi}$ to be a 1-step forward model $\hat{z}_{i}=\hat{z}_{t+1}=g_{\phi}(z_{t},a_{t})$. We set our positive labels $y^{\prime}$ to be the true representation of the next frame $z^{\prime}_{i}=z^{\prime}_{t+1}=f_{\theta}(o(x_{t+1},e^{\prime}))$.

We focus on the following environments (see Fig. 2).

To evaluate sim-to-real transfer, we use the Franka Emika Panda robotic arm, with an Intel RealSense D415 camera positioned to look diagonally down at the scene (see Fig. 5). We make sure that the arm is not visible in the images by moving it out of sight before taking an image. We also use a simulation of the real Panda robotic arm environment to generate training data and for evaluation. Our manipulation task is moving a random sized cube. The actions are similar to the Controlled Deformable Rope environment – the Panda robotic arm pushes the cube in a random direction that is parallel to the table, specified by a start and goal position for the tip of the gripper.

Our goal is to learn representations that capture relevant properties of the scene, and ignore spurious features. To evaluate such behavior on out-of-distribution data, we test the ability to retrieve similar observations that share similar physical properties (e.g. position, orientation, size) but from different domains. Given a test image $o$, we first encode it into $z=f_{\theta}(o)$, and retrieve the latent nearest neighbour from an external dataset $D_{ext}$, unseen during training $z_{NN}=\operatorname*{arg\,min}_{z^{\prime}\in D_{ext}}h(z,z^{\prime})$. For $h$, we use the Cosine distance, as we found it to perform well for both CDR and the baseline. We evaluate test images from both in distribution (following the terminology of Sec. III-B, with $e$ seen in training but different $x$) and OOD (both $e$ and $x$ were not seen in training). To evaluate performance, we measure the Intersection Over Union (IOU) of the object pixels between the nearest neighbor and the test image (higher is better). For the rigid object environments, we also report the sum of Euclidean distances between the objects in the test image and its retrieved nearest neighbour (lower is better). Table I shows quantitative results of CDR against Baseline. Note that CDR significantly outperforms the baselines by a large margin.

To test if our representations are in fact domain-invariant, we randomly sample observations from different domains but with the same object configurations $o_{1}=o(x,e_{1}),o_{2}=o(x,e_{2})$, and compare the distance between them in the latent space, i.e $h(o_{1},o_{2})$. Table II shows that CDR is significantly more domain-invariant compare to the baseline.

To evaluate sim-to-real transfer, we use real world images that are very different from the simulated data in both domain appearance and physical properties (see Fig. 3 and Fig. 4). We then qualitatively evaluate our ability to retrieve similar simulated data. In Fig. 3 we compare CDR and Baseline to retrieve observations from a challenging reference image with a very different lightning and frame. Though retrieved positions are not always perfect, CDR tends to capture correctly the relative position and sizes of the objects – an essential property for predicting the outcome of a collision. For the controlled rigid objects data, we use a hammer and a banana much bigger than the maximal size in simulation. Nevertheless, CDR tends to capture correctly the relationship between the objects and the cube agent. For the rope data, we use a colored rope, and unlike previous work [11], we use a long rope that can form complex shapes. Furthermore, to demonstrate generalization outside lab conditions, we also use the same rope material as background for some reference images. Note that CDR yields retrievals that are much more accurate than the baseline.

To evaluate if representations learned with CDR are useful for planning, we test our model in reaching a goal image from a random state using a simple 1-step Model Predictive Control (MPC) planner. The planner is provided with a goal image $o_{goal}$ which can be either a real-world image or a simulated one (see Fig. 6), which is then encoded into $z_{goal}=f_{\theta}(o_{goal})$. In each time step, the planner samples a set $A$ of 1000 possible actions that are guaranteed to push the cube in some random direction and magnitude. We feed these into our forward model, along with the representation of the current image $z_{t}=f_{\theta}(o_{t})$ to get a set of 1000 latent representations of next states, corresponding to the different actions, $\hat{z}_{t+1}=g_{\phi}(z_{t},a)$. We greedily choose an action to execute, based on distance in latent space $a_{t}=\operatorname*{arg\,min}_{a\in A}h(z_{goal},g_{\phi}(z_{t},a))$.

To evaluate sim2real transfer, we use both simulated goal images and real-world goal images, and test the robot on a variety of challenging environments. For example, we add multiple objects to the scene, while no additional objects were seen during training (see Fig. 7). We measure the Euclidean distance from the final state to the goal state. Table III shows that CDR significantly outperforms the CFM Baseline.