# An Empirical Investigation of
Representation Learning for Imitation

Xin Chen
The University of Hong Kong
<EMAIL_ADDRESS>
&Sam Toyer11footnotemark: 1
UC Berkeley
<EMAIL_ADDRESS>
&Cody Wild11footnotemark: 1
UC Berkeley
<EMAIL_ADDRESS>
&Scott Emmons
UC Berkeley
&Ian Fischer
Google Research
&Kuang-Huei Lee
Google Research
&Neel Alex
UC Berkeley
&Steven Wang
UC Berkeley
&Ping Luo
The University of Hong Kong
&Stuart Russell
UC Berkeley
&Pieter Abbeel
UC Berkeley
&Rohin Shah
UC Berkeley
Equal contribution, corresponding authors

###### Abstract

Imitation learning often needs a large demonstration set in order to handle
the full range of situations that an agent might find itself in during
deployment. However, collecting expert demonstrations can be expensive. Recent
work in vision, reinforcement learning, and NLP has shown that auxiliary
representation learning objectives can reduce the need for large amounts of
expensive, task-specific data. Our Empirical Investigation of Representation
Learning for Imitation (EIRLI) investigates whether similar benefits apply to
imitation learning. We propose a modular framework for constructing
representation learning algorithms, then use our framework to evaluate the
utility of representation learning for imitation across several environment
suites. In the settings we evaluate, we find that existing algorithms for
image-based representation learning provide limited value relative to a well-
tuned baseline with image augmentations. To explain this result, we
investigate differences between imitation learning and other settings where
representation learning has provided significant benefit, such as image
classification. Finally, we release a well-documented codebase which both
replicates our findings and provides a modular framework for creating new
representation learning algorithms out of reusable components.

##  1 Introduction

Much recent work has focused on how AI systems can learn what to do from human
feedback [1]. The most popular approach—and the focus of this paper—is
imitation learning (IL), in which an agent learns to complete a task by
mimicking demonstrations of a human.

As demonstrations can be costly to collect, we would like to learn
representations that lead to better imitation performance given limited data.
Many existing representation learning (RepL) methods in Computer Vision and
Reinforcement Learning do exactly this, by extracting effective visual [2] or
temporal [3] information from inputs. A natural hypothesis is that RepL would
also add value for IL.

We test this hypothesis by investigating the impact of common RepL algorithms
on Behavioral Cloning (BC) and Generative Adversarial Imitation Learning
(GAIL). We survey a wide variety of RepL methods, and construct a modular
framework in which each design decision can be varied independently. As
previous work has found that image augmentation alone can outperform more
complex representation learning techniques [4, 5], we make sure to compare
against baselines that use augmentation. To ensure generalizability of our
results, we evaluate on ten tasks selected across three benchmarks, including
MAGICAL [6], Procgen [7] and the DeepMind Control Suite (DMC) [8].

We find that, on average, RepL methods do significantly outperform vanilla BC,
but this benefit can be obtained simply by applying well-tuned image
augmentations during BC training. To understand the discrepancy between this
result and the success of RepL in computer vision and reinforcement learning,
we apply clustering algorithms and attribution methods to qualitatively
investigate the learned representations and policies, surfacing a number of
intriguing hypotheses for investigation in future work.

This paper is, to the best of our knowledge, the first to provide a systematic
empirical analysis of different representation learning methods for imitation
learning in image-based environments. Concretely, our Empirical Investigation
of Representation Learning for Imitation (EIRLI) makes the following
contributions:

  1. 1.

We identify meaningful axes of variation in representation learning algorithm
design, allowing us to construct a modular framework to conceptually analyze
these designs.

  2. 2.

We use this framework to build a well documented, modular, and extensible code
base, which we release at github.com/HumanCompatibleAI/eirli.

  3. 3.

We conduct an extensive comparison of popular RepL methods in the imitation
learning setting, and show that RepL has limited impact on task performance
relative to ordinary image augmentations. By analysing our learned
representations and policies, we identify several promising directions for
future work at the intersection of representation learning and decision-
making.

##  2 Design decisions in representation learning

Figure 1: A framework for the use of representation learning (RepL) in
imitation learning. In the pretraining setting, we first train the encoder
with RepL, then finetune end-to-end with IL. In the joint training setting,
the RepL objective is used as an auxiliary loss throughout IL training.

To apply representation learning (RepL) effectively, it is important to
understand the relative impact of different RepL algorithm design choices on
downstream task performance. We argue that for many common RepL algorithms,
these design choices can be broken down along a common set of axes, which we
show in Table 1 and Table 2. In this section, we elaborate on our conceptual
breakdown both as a literature review and as an implementation walkthrough of
our RepL framework.

We summarize existing RepL for image classification algorithms in Table 1 and
a selection of RepL for reinforcement learning algorithms in Table 2. The full
version of the table deconstructing current RepL methods in reinforcement
learning can be found in the appendix in Table A1. [5]

Table 1: Design choices made in representation learning for image recognition. “Augmentation”, “Momentum”, and “Projection” show whether image augmentation, target encoder momentum, and projection heads were used, respectively. “Pre/Joint” shows whether RepL is used as a pretraining step, or is jointly learned with the downstream task (typically as an auxiliary loss). Algorithm | Task | Augmentation | Momentum | Projection | Pre/Joint
---|---|---|---|---|---
VAE [9] | Reconstruction | ✗ | ✗ | ✗ | Pre
AugMix [10] | Consistency | ✓ | ✗ | ✗ | Joint
FixMatch [11] | Consistency | ✓ | ✗ | ✗ | Joint
CPC [12] | Contrastive | ✗ | ✗ | ✓ | Pre
MoCo [13] | Contrastive | ✓ | ✓ | ✗ | Pre
SimCLR [2] | Contrastive | ✓ | ✗ | ✓ | Pre
SimCLRv2 [14] | Contrastive | ✓ | ✗ | ✓ | Pre
BYOL [15] | Bootstrap | ✓ | ✓ | ✓ | Pre
Table 2: Design choices made in a selection of representation learning algorithms for reinforcement learning (full table in the appendix). Act, Aug, Mom, Proj and Comp respectively show whether action conditioning, augmentation, momentum, projection heads, and compression were used. P/J determines whether the representation learning is an initial (P)retraining step, or is (J)ointly learned alongside reinforcement learning. R/C in the Task column refer to Reconstruction/Contrastive. Note that different papers may use different sets of augmentations. Algorithm | Task | RL alg. | Context | Target | Act | Aug | Mom | Proj | Comp | P/J
---|---|---|---|---|---|---|---|---|---|---
World models [16] | R | CMA-ES | $o_{t}$ | $o_{t},o_{t+1}$ | ✓ | ✗ | ✗ | ✗ | ✗ | P
PlaNet [17] | R | MPC + CEM | $o_{1:t}$ | $o_{t+1:T},r_{t+1:T}$ | ✓ | ✗ | ✗ | ✗ | ✗ | J
CURL [18] | C | SAC | $o_{t}$ | $o_{t}$ | ✗ | ✓ | ✗ | ✗ | ✗ | J
PI-SAC [3] | C | SAC | $o_{t}$ | $o_{t+k},r_{t+k}$ | ✓ | ✓ | ✓ | ✓ | ✓ | J
ATC [19] | C | SAC, PPO | $o_{t}$ | $o_{t+k}$ | ✗ | ✓ | ✓ | ✓ | ✗ | P

###  2.1 Target selection

Most RepL methods can be thought of as proxy tasks in which a dataset of
$(x,y)$ pairs is provided and the network must model some aspects of the
relationship between $x$ and $y$. Since the learning signal derives from the
relationship between $x$ and $y$, the choice of $x$ and $y$ thus has a
significant impact on exactly what information is modeled. We refer to the
inputs for which representations $z$ are computed as the “context” $x$, and
the inputs with which they are related are the “targets” $y$. Often, the
target is a (possibly transformed) version of a context.

In image classification, learned representations must capture the label-
relevant information in a single input image. It is assumed that most images
used for representation learning will not have labels or other task-relevant
metadata. Thus, the context and target are typically both set to the original
image, after which they may be augmented in different ways. For example, in a
Variational Autoencoder (VAE) [9], an input image (context) is encoded into a
vector representation and then decoded back into pixels, which is then
compared against the same input image (now interpreted as a target).

Once we move to sequential decision-making, the observations have a sequential
structure, and there is a notion of actions and a reward function. These can
all be leveraged in the construction of the contexts and targets. For example,
a _Temporal VAE_ is identical to a regular VAE, except that for a context
observation $o_{t}$, we set the target to be a future observation $o_{t+k}$.
Now, the input observation $o_{t}$ (context) is encoded into a vector
representation and then decoded back into pixels, which is then compared
against the _future_ observation $o_{t+k}$ (target). By using a temporal
target, we now incentivize representations that contain _predictive_
information [20]. In reinforcement learning, another option is to add the
reward $r_{t}$ to the target to encourage learning representations that are
useful for planning.

###  2.2 Loss type

We divide modern methods for representation learning into four categories:

Reconstruction. Here, the goal is to reconstruct the target $y$ from the
representation $z$. Both the VAE and temporal VAE in the previous section use
a reconstructive loss, in which a _decoded image_ $d_{\phi}(z)$ is compared
against the target $y$, and that reconstruction loss is combined with a
regularization term.

Contrast. Contrastive methods take a series of context–target pairs
$(x_{1},y_{1}),(x_{2},y_{2}),\ldots,(x_{K},y_{K})$ and use the same network to
encode both the context and target into latent representations $z_{i}\sim
e(z\mid x_{i})$ and $z_{i}^{\prime}\sim e(z\mid y_{i})$. A contrastive loss
then incentivizes $z_{i}$ and $z_{i}^{\prime}$ to be similar to each other,
but different from $z_{j}$ and $z_{j}^{\prime}$ for all other pairs $j\neq i$.
Typically, the contrastive loss function is chosen to maximize the mutual
information $I(z;y)$, such as with the InfoNCE loss function [21]:

$\mathcal{L}_{\text{InfoNCE}}=\mathbb{E}\left[\log\frac{e^{f\left(x_{i},y_{i}\right)}}{\frac{1}{K}\sum_{j=1}^{K}e^{f\left(x_{i},y_{j}\right)}}\right]$

$f$ could, for instance, be a bilinear function
$f(x_{i},y_{i})=z_{i}^{T}Wz_{i}^{\prime}$, where $z_{i}\sim e(z\mid x_{i})$,
$z_{i}^{\prime}\sim e(z\mid y_{i})$, and $W\in\mathbb{R}^{n\times n}$ is a
learned parameter matrix.

Bootstrapping. This is a simplified variant of contrastive learning. Given a
related context $x$ and target $y$, a bootstrapping method predicts a moving-
average-encoded target from the encoded context. Bootstrapping does not need a
large dataset of negatives to prevent the representation from collapsing to a
single point; instead, it prevents collapse by stopping gradients from
propagating through the target encoder.

Consistency. These methods, such as AugMix [10] and FixMatch [11], include
auxiliary loss terms that encourage the model to produce similar distributions
over $y$ for different transformations of the same input image.

Compression. A representation $z\sim e_{\theta}(\cdot\mid x)$ should contain
enough information about the input $x$ to solve downstream tasks. Ideally,
$e_{\theta}$ should also extract only the minimum amount of information about
$x$ that is necessary to perform well. We refer to this as compression. As a
form of explicit compression, we implement the _conditional entropy
bottleneck_ (CEB) [22], which approximately minimizes $I(X;Z\mid Y)$.

###  2.3 Augmentation

In many algorithm designs, one or both of the context frame and target frame
undergo augmentation before being processed by the encoder and decoder
networks. In some algorithms, like SimCLR, this augmentation is the main
source of noise causing transformed representations of the same input to not
be purely identical. In other algorithms, it simply helps promote
generalization by sampling from a wider image distribution than would be done
naturally.

###  2.4 Neural network

In the case of a VAE, the neural network consists of two parts. The _encoder_
produces the latent representation from the input, while the _decoder_
reconstructs the input from the latent representation. We generalize this
terminology and _define_ the encoder for an arbitrary RepL method to be that
part of the neural network that is used to compute the representation, and the
decoder to be the rest of the neural network. Under this definition, the
downstream tasks (which could include imitation, classification, reinforcement
learning, etc.) only require the encoder, not the decoder. Note that the
“decoder” may not convert the learned representation into some human-
interpretable format; it is simply those parts of the neural network that are
required by the RepL method but that do not serve to compute the
representation.

####  2.4.1 Encoder

The encoder is the core component of a representation learner: it is
responsible for mapping input targets $x$ into $z$ vectors that are used as
the learnt representation in downstream tasks.

Recurrent encoders. In some cases, a “context” could be a sequence of frames
instead of a single frame, and the encoder could compress that into a single
representation of the past. This paper doesn’t address recurrent encoders,
opting instead to make all encoders operate on single framestacks.

Momentum encoders. In contrastive tasks, learning a high-quality
representation often requires large batch sizes, since the difficulty of the
contrastive task scales with the number of negatives. However, batches of the
appropriate difficulty can be so large that encoding the negative targets
becomes prohibitively compute- and memory-intensive. He et al. [13] propose
reusing negative targets from previous batches to alleviate this cost. One
challenge with reusing targets is that the encoder can change too quickly
during training, in which case negative targets from previous batches become
“stale”. Thus He et al. [13] use a separate target encoder which is updated
slowly enough that targets do not become stale too quickly. Specifically, the
target encoder’s weights $\theta_{t}$ are updated to track the main context
encoder weights $\theta_{c}$ using the update rule
$\theta_{t}\leftarrow\alpha\theta_{t}+(1-\alpha)\theta_{c}$. $\alpha$ is
referred to as a momentum parameter, and is typically set to some value close
to 1 (e.g. $\alpha=0.999$).

####  2.4.2 Decoder

Decoders are optional neural network layers applied before a loss is
calculated, but which are _not_ included in the learnt encoder used at
transfer time. They take in the $z$ output by the encoder (and optional
additional information), and produce an input to the loss function.

Image reconstruction. The most common historical form of decoder in a RepL
algorithm is the image reconstruction decoder, which has historically been
used by VAEs and similar model designs to “decode” a predicted image from a
representation bottleneck. This predicted image is used in calculating a MLE
loss against the true image, but is discarded before downstream transfer
tasks.

Projection heads. Projection heads are multi-layer perceptrons that take in
the output of the encoder and project it into a new space over which the loss
can then be calculated. Recent work has shown these to be useful for
contrastive learning [2].

Action conditioning. Temporal tasks can be made easier by conditioning on the
action $a_{t}$. However, for an encoder to be used for reinforcement learning
or imitation, the representation must not depend on the current action
$a_{t}$. Thus, the encoder is only responsible for learning a $z$
representation of the observation $o_{t}$, and is combined with a
representation of the action within the decoder step.

###  2.5 Pretraining vs joint training

Another question is how to integrate representation learning with an RL
algorithm. In image recognition, representation learning is done as a
pretraining step. We experiment with this approach in this work, as well as
the strategy of "joint training", where we add the representation learning
loss as an _auxiliary loss_ while performing reinforcement learning.

##  3 Experiments

Given our framework, it is straightforward to construct RepL algorithms that
differ along any of the axes described in Section 2. In this section, we
create a representative set of such algorithms and evaluate various ways of
combining them with imitation learning. Although some RepL methods appear to
be effective on some tasks, we find that the difference between using and not
using RepL is often much less than the difference between using and not using
augmentations for the imitation policy. In Section 4, we discuss possible
reasons why RepL does not have a greater effect, and suggest alternative ways
that RepL could be used more fruitfully.

###  3.1 Experiment setup

Environments and training data. We evaluate on ten tasks taken from three
benchmark domains: DMC [8], Procgen [7], and MAGICAL [6]. Here we briefly
explain our choice of tasks and datasets; for more detailed information (e.g.
dataset sizes and collection methods), refer to Appendix D.

From DMC, we take image-based versions of the cheetah-run, finger-spin, and
reacher-easy tasks. All three of these are popular benchmark tasks for deep RL
and deep IL, and represent a range of difficulties (reacher-easy being the
easiest, and cheetah-run being the hardest). However, they provide limited
evaluation of generalisation. We use a common demonstration set for RepL and
IL.

Table 3: Design decisions for representation learning algorithms used in our experiments. Algorithm | Task | Context | Target | Act | Aug
---|---|---|---|---|---
Temporal CPC | Contrastive | $o_{t}$ | $o_{t+1}$ | ✗ | ✓
SimCLR | Contrastive | $o_{t}$ | $o_{t}$ | ✗ | ✓
VAE | Reconstructive | $o_{t}$ | $o_{t}$ | ✗ | ✗
Dynamics | Reconstructive | $o_{t},a_{t}$ | $o_{t+1}$ | ✓ | ✗
Inverse Dynamics | Reconstructive | $o_{t},o_{t+1}$ | $a_{t}$ | ✗ | ✗

From Procgen, we choose the “easy” variants of the CoinRun, Fruitbot, Jumper
and Miner tasks. In Procgen, different random initialisations for a given task
can have wildly different appearance and structure, but still admit a common
optimal policy. This makes it a much more challenging evaluation of
generalization than DMC. As with DMC, we use the same demonstration set for
RepL and IL.

From MAGICAL we choose the MoveToRegion, MoveToCorner, and MatchRegions tasks,
which represent a range of difficulty levels (MoveToRegion being the easiest,
and MatchRegions being the hardest). For each task, MAGICAL defines a “demo
variant” for training and a set of “test variants” for evaluating robustness
to changes in dynamics, appearance, etc. Unlike DMC and Procgen, our MAGICAL
experiments augment the demonstration set with additional demo variant random
rollouts for RepL training. This models the setting in which it is cheap to
collect additional data for self-supervised learning, but expensive to collect
demonstrations. We include more detailed environment setups in Appendix B.

Imitation baselines. Most of our experiments use behavioral cloning (BC) [23]
as the base imitation learning algorithm. Given a dataset
$\mathcal{D}=\\{(x_{0},a_{0}),(x_{1},a_{1}),\ldots\\}$ of observation–action
tuples drawn from a demonstrator, BC optimises the policy $\pi_{\theta}(a\mid
x)$ to maximise the expected log likelihood,

$\mathcal{L}(\theta)=\mathop{\mathbb{E}}_{(x,a)\sim\mathcal{D}}\left[{\log\pi_{\theta}(a\mid
x)}\right]\,.$

We combine BC with representation learning in two ways. First, we use RepL to
pretrain all but the final layer of the policy, then fine-tune the policy end-
to-end with BC. This appears to be the most popular approach in the vision
literature. Second, we use RepL as an auxiliary objective during BC training,
so that both imitation and representation learning are performed
simultaneously. Importantly, we also do control runs both with and without
image augmentations. The deep RL community has repeatedly found that image
augmentations can yield a greater improvement than some sophisticated
representation learning methods [18, 5], and so it is important to distinguish
between performance gains due to the choice of RepL objective and performance
gains due to the use of augmentations.

In addition to BC, we present results with Generative Adversarial Imitation
Learning (GAIL) [24] and RepL pretraining. GAIL treats IL as a game between an
imitation policy $\pi_{\theta}(a\mid x)$ and a discriminator $D_{\psi}(x,a)$
that must distinguish $\pi_{\theta}$’s behaviour from that of the
demonstrator. Using alternating gradient descent, GAIL attempts to find a
$\theta$ and $\psi$ that attain the saddle point of

$\displaystyle\max_{\theta}\min_{\psi}\left\\{-\mathop{\mathbb{E}}_{(x,a)\sim\pi_{\theta}}\left[{\log
D_{\psi}(x,a)}\right]-\mathop{\mathbb{E}}_{(x,a)\sim\mathcal{D}}\left[{\log(1-D_{\psi}(x,a))}\right]+w_{H}H(\pi_{\theta})\right\\}\,.$

Here $H$ is an entropy penalty weighted by regularisation parameter $w_{H}\geq
0$. We use augmentations only for the GAIL discriminator, and not the policy
(we could not get GAIL to train reliably with policy augmentations).
Discriminator regularisation is of particular interest because past work has
shown that discriminator augmentations are essential to obtaining reasonable
imitation performance when applying GAIL to image-based environments [25]. For
our experiments combining GAIL with RepL, we use the learned representation to
initialize both the GAIL discriminator and the GAIL policy.

RepL algorithms. Using our modular representation learning framework, we
construct five representation learning algorithms described in Table 3. More
detailed descriptions are in Appendix B.

###  3.2 Results

Results are shown in Table 4 for BC + RepL pretraining, and Table 5 for BC +
RepL joint training, and Table 6 for GAIL + RepL pretraining. Each cell shows
mean ${\pm}$ standard deviation over at least five random seeds. We treat IL
with augmentations (but no RepL) as our baseline. We color cells that have a
higher mean return than the baseline, and mark them with an asterisk (*) when
the difference is significant at $p<0.05$, as measured by a one-sided Welch’s
t-test without adjustment for multiple comparisons. We include the loss curves
for our BC experiments in Appendix H.

Table 4: Pretraining results for BC. We color cells that have a higher mean return than BC with augmentations, and mark them with an asterisk (*) when the difference is significant at $p<0.05$, as measured by a one-sided Welch’s t-test without adjustment for multiple comparisons. Env | Task | Dynamics | InvDyn | SimCLR | TemporalCPC | VAE | BC aug | BC no aug |
---|---|---|---|---|---|---|---|---|---
DMC | cheetah-run | 482${\pm}$36 | 669${\pm}$18 | 687${\pm}$17 | 661${\pm}$13 | 458${\pm}$39 | 690${\pm}$17 | 617${\pm}$34 |
| finger-spin | 718${\pm}$17 | 748${\pm}$17* | 726${\pm}$1 | 723${\pm}$4 | 751${\pm}$6* | 730${\pm}$9 | 940${\pm}$4* |
| reacher-easy | 774${\pm}$24 | 890${\pm}$14 | 907${\pm}$9 | 893${\pm}$13 | 880${\pm}$20 | 874${\pm}$21 | 452${\pm}$34 |
Procgen | coinrun-train | 8.1${\pm}$0.4 | 8.0${\pm}$0.2 | 8.0${\pm}$0.5 | 8.1${\pm}$0.3 | 8.4${\pm}$0.4 | 8.1${\pm}$0.3 | 8.7${\pm}$0.6* |
| fruitbot-train | 3.2${\pm}$1 | 16.2${\pm}$1.2 | 17.5${\pm}$1.9 | 15.4${\pm}$1.5 | 17.5${\pm}$1.5 | 18.3${\pm}$1.9 | 11.4${\pm}$0.6 |
| jumper-train | 8.1${\pm}$0.2 | 8.0${\pm}$0.4 | 7.9${\pm}$0.6 | 7.5${\pm}$0.6 | 7.9${\pm}$0.6 | 8.1${\pm}$1.2 | 7.1${\pm}$1.2 |
| miner-train | 4.5${\pm}$1.2 | 5.9${\pm}$0.2 | 9.9${\pm}$0.4 | 9.5${\pm}$2.3 | 10.4${\pm}$0.3* | 9.8${\pm}$0.3 | 8.1${\pm}$0.3 |
| coinrun-test | 6.3${\pm}$0.8 | 6.9${\pm}$0.5 | 6.8${\pm}$0.5 | 6.8${\pm}$0.4 | 7.0${\pm}$0.5 | 6.7${\pm}$0.4 | 6.5${\pm}$0.7 |
| fruitbot-test | -3${\pm}$0.9 | 15.6${\pm}$1.1 | 13.4${\pm}$1.0 | 14.7${\pm}$1.0 | 13.2${\pm}$1.0 | 13.7${\pm}$1.1 | 2.2${\pm}$0.6 |
| jumper-test | 3.2${\pm}$0.4 | 3.9${\pm}$0.3 | 3.6${\pm}$0.4 | 3.7${\pm}$0.5 | 3.4${\pm}$0.5 | 3.9${\pm}$0.5 | 4.6${\pm}$0.4 |
| miner-test | 0.6${\pm}$0.1 | 2.6${\pm}$0.1 | 2.6${\pm}$0.4 | 3.1${\pm}$0.4 | 2.7${\pm}$0.3 | 2.7${\pm}$0.4 | 0.8${\pm}$0.1 |
MAGI- | MatchRegions | 0.42${\pm}$0.04 | 0.42${\pm}$0.04 | 0.42${\pm}$0.03 | 0.41${\pm}$0.01 | 0.42${\pm}$0.03 | 0.43${\pm}$0.02 | 0.28${\pm}$0.08 |
CAL | MoveToCorner | 0.84${\pm}$0.07 | 0.83${\pm}$0.04 | 0.83${\pm}$0.04* | 0.80${\pm}$0.02 | 0.78${\pm}$0.06 | 0.78${\pm}$0.05 | 0.72${\pm}$0.04 |
| MoveToRegion | 0.82${\pm}$0.02* | 0.83${\pm}$0.02* | 0.82${\pm}$0.01* | 0.81${\pm}$0.01* | 0.81${\pm}$0.05* | 0.74${\pm}$0.02 | 0.81${\pm}$0.04* |

BC pretraining results. In the pretraining setting, we see that none of our
RepL algorithms consistently yield large improvements across all (or even
most) tasks. Indeed, the relative impact of adding representation learning
tends to be lower than the impact of adding or removing augmentations.
Although adding augmentations to BC usually yields a large improvement, there
are a handful of tasks where adding augmentations substantially decreases
performance; we remark further on this below. Note that most of our RepL
algorithms do seem to yield an improvement in MoveToRegion, suggesting that
there may still be value to RepL for a narrower set of tasks and datasets.

BC joint training results. When using joint training as an auxiliary loss, we
similarly see that no one RepL method consistently improves performance across
all benchmark tasks. However, in the DMC tasks, we do see consistent
improvement over the baseline for all RepL methods. This suggests that our
RepL methods provide benefit in some environments, but are sensitive to the
choice of task.

Table 5: Joint training results for BC. We color cells that have a higher mean return than BC with augmentations, and mark them with an asterisk (*) when the difference is significant at $p<0.05$, as measured by a one-sided Welch’s t-test without adjustment for multiple comparisons. Env | Task | Dynamics | InvDyn | SimCLR | TemporalCPC | VAE | BC aug | BC no aug |
---|---|---|---|---|---|---|---|---|---
DMC | cheetah-run | 723${\pm}$14* | 716${\pm}$23* | 717${\pm}$11* | 716${\pm}$16* | 724${\pm}$12* | 690${\pm}$17 | 617${\pm}$34 |
| finger-spin | 755${\pm}$6* | 755${\pm}$12* | 732${\pm}$15 | 725${\pm}$12 | 755${\pm}$3* | 730${\pm}$9 | 940${\pm}$4* |
| reacher-easy | 898${\pm}$19 | 903${\pm}$10* | 889${\pm}$14 | 912${\pm}$18* | 903${\pm}$8* | 874${\pm}$21 | 452${\pm}$34 |
Proc- | coinrun-train | 8.0${\pm}$0.4 | 7.1${\pm}$0.3 | 8.0${\pm}$0.5 | 8.6${\pm}$0.5* | 7.9${\pm}$0.2 | 8.1${\pm}$0.3 | 8.7${\pm}$0.6* |
gen | fruitbot-train | 17.0${\pm}$0.7 | 6.6${\pm}$1.4 | 13.4${\pm}$1.9 | 11.4${\pm}$0.7 | 15.4${\pm}$1.0 | 18.3${\pm}$1.9 | 11.4${\pm}$0.6 |
| jumper-train | 7.9${\pm}$0.5 | 8.1${\pm}$0.4 | 8.0${\pm}$0.4 | 8.0${\pm}$0.3 | 8.3${\pm}$0.5 | 8.1${\pm}$1.2 | 7.1${\pm}$1.2 |
| miner-train | 8.9${\pm}$0.8 | 8.9${\pm}$0.7 | 8.7${\pm}$0.3 | 7.1${\pm}$0.8 | 8.6${\pm}$0.7 | 9.8${\pm}$0.3 | 8.1${\pm}$0.3 |
| coinrun-test | 6.4${\pm}$0.4 | 6.0${\pm}$0.5 | 6.6${\pm}$0.3 | 6.2${\pm}$0.5 | 6.9${\pm}$0.4 | 6.7${\pm}$0.4 | 6.5${\pm}$0.7 |
| fruitbot-test | 10.9${\pm}$0.7 | 3.3${\pm}$1.1 | 8.5${\pm}$1.5 | 6.4${\pm}$1.2 | 10.4${\pm}$1.6 | 13.7${\pm}$1.1 | 2.2${\pm}$0.6 |
| jumper-test | 3.4${\pm}$0.3 | 4.8${\pm}$0.2* | 3.8${\pm}$0.3 | 3.4${\pm}$0.3 | 3.9${\pm}$0.7 | 3.9${\pm}$0.5 | 4.6${\pm}$0.4* |
| miner-test | 2.0${\pm}$0.2 | 1.9${\pm}$0.3 | 1.8${\pm}$0.3 | 1.0${\pm}$0.2 | 2.0${\pm}$0.3 | 2.7${\pm}$0.4 | 0.8${\pm}$0.1 |
MAGI- | MatchRegions | 0.44${\pm}$0.02 | 0.23${\pm}$0.08 | 0.41${\pm}$0.02 | 0.01${\pm}$0.01 | 0.41${\pm}$0.03 | 0.43${\pm}$0.03 | 0.31${\pm}$0.02 |
CAL | MoveToCorner | 0.78${\pm}$0.07 | 0.30${\pm}$0.22 | 0.76${\pm}$0.05 | 0.02${\pm}$0.02 | 0.82${\pm}$0.06 | 0.80${\pm}$0.05 | 0.70${\pm}$0.09 |
| MoveToRegion | 0.76${\pm}$0.02 | 0.35${\pm}$0.24 | 0.74${\pm}$0.01 | 0.47${\pm}$0.07 | 0.77${\pm}$0.02 | 0.75${\pm}$0.02 | 0.78${\pm}$0.04 |

Effect of augmentations on BC. Incorporating augmentations into BC training
tended to yield the largest effect of any technique considered in this work,
even without an explicit representation learning loss. In roughly half of the
environments studied, this had a substantial impact on reward, and reward
increased $150\%$ or more in reacher-easy, Fruitbot, and MatchRegions.
However, environments seem to be bimodal in their response to augmentations:
in a handful of environments (finger-spin, coinrun-train, jumper-test, and
MoveToRegion), adding augmentations leads to consistently _worse_ performance.
This effect is particularly dramatic in finger-spin, which we believe is a
result of the fact that relevant objects in the environment always stay fixed.
Consequently, translational augmentations don’t aid generalization, and
rotational augmentations can be confused with true signal (since the angle of
the finger determines the ideal action). Because augmentation already yields
large benefits, many of the representation learning algorithms do not provide
much additional gain on top of BC-Augs, even when they perform substantially
better than BC-NoAugs. This result is consistent with the finding by Laskin et
al. [4] that simply augmenting input frames in reinforcement learning produced
performance on par with sophisticated representation learning methods.

GAIL pretraining results. GAIL pretraining results mirror those for BC
pretraining, but with even fewer statistically significant deviations from
baseline performance. We see that augmentation can be even more important for
GAIL than it is for BC. For instance, GAIL with discriminator augmentations
obtains higher return on finger-spin than BC does, but obtains a return of 0
when discriminator augmentations are removed. This is consistent with the
observation of Zolna et al. that strict regularisation is essential to make
GAIL perform well in image-based domains [25].

Table 6: Pretraining results for GAIL. We color cells that have a higher mean return than BC with augmentations, and mark them with an asterisk (*) when the difference is significant at $p<0.05$, as measured by a one-sided Welch’s t-test without adjustment for multiple comparisons. For the sake of space, we abbreviate TemporalCPC to $t$CPC. Env | Task | Dynamics | InvDyn | SimCLR |  $t$CPC | VAE | GAIL aug | GAIL no aug
---|---|---|---|---|---|---|---|---
DMC | cheetah-run | 380${\pm}$76 | 320${\pm}$61 | 265${\pm}$58 | 360${\pm}$74 | 375${\pm}$33 | 449${\pm}$67 | 75${\pm}$40
| finger-spin | 868${\pm}$14 | 886${\pm}$8* | 800${\pm}$23 | 748${\pm}$72 | 868${\pm}$18 | 868${\pm}$12 | 0${\pm}$0
| reacher-easy | 53${\pm}$24 | 73${\pm}$51 | 21${\pm}$23 | 118${\pm}$88 | 122${\pm}$89 | 221${\pm}$162 | 89${\pm}$88
Proc- | coinrun-train | 5.9${\pm}$0.29* | 5.85${\pm}$0.51* | 2.15${\pm}$1.53 | 3.28${\pm}$2.62 | 3.54${\pm}$1.22 | 3.31${\pm}$0.44 | 2.80${\pm}$0.89
gen | fruitbot-train | -2.81${\pm}$0.1 | -2.37${\pm}$0.55 | -2.47${\pm}$0.15 | -2.38${\pm}$0.31 | -2.49${\pm}$0.22 | -2.42${\pm}$0.42 | -2.63${\pm}$0.30
| jumper-train | 3.31${\pm}$0.31 | 3.17${\pm}$0.40 | 3.36${\pm}$0.53 | 2.69${\pm}$1.31 | 3.45${\pm}$0.70 | 3.44${\pm}$0.52 | 3.47${\pm}$0.53
| miner-train | 0.53${\pm}$0.12 | 0.60${\pm}$0.11 | 0.53${\pm}$0.14 | 0.84${\pm}$0.14* | 0.51${\pm}$0.07 | 0.65${\pm}$0.10 | 0.77${\pm}$0.18
| coinrun-test | 6.1${\pm}$0.9* | 5.91${\pm}$0.16* | 2.11${\pm}$1.61 | 3.35${\pm}$2.74 | 3.01${\pm}$1.10 | 3.44${\pm}$0.68 | 2.77${\pm}$0.84
| fruitbot-test | -2.44${\pm}$0.49 | -2.65${\pm}$0.24 | -2.55${\pm}$0.30 | -2.65${\pm}$0.14 | -2.85${\pm}$0.33 | -2.44${\pm}$0.50 | -2.51${\pm}$0.44
| jumper-test | 2.56${\pm}$0.52 | 2.53${\pm}$0.64 | 3.15${\pm}$0.45 | 2.35${\pm}$0.81 | 2.75${\pm}$0.59 | 3.25${\pm}$0.42 | 3.15${\pm}$0.20
| miner-test | 0.36${\pm}$0.04 | 0.57${\pm}$0.07 | 0.55${\pm}$0.24 | 0.87${\pm}$0.15* | 0.50${\pm}$0.17 | 0.65${\pm}$0.17 | 0.66${\pm}$0.14
MAGI- | MatchRegions | 0.42${\pm}$0.10 | 0.34${\pm}$0.12 | 0.47${\pm}$0.04 | 0.39${\pm}$0.12 | 0.30${\pm}$0.15 | 0.46${\pm}$0.06 | 0.22${\pm}$0.12
CAL | MoveToCorner | 0.48${\pm}$0.09 | 0.45${\pm}$0.10 | 0.52${\pm}$0.07 | 0.55${\pm}$0.15 | 0.62${\pm}$0.11* | 0.49${\pm}$0.08 | 0.55${\pm}$0.14
| MoveToRegion | 0.72${\pm}$0.07 | 0.74${\pm}$0.04 | 0.74${\pm}$0.06 | 0.76${\pm}$0.03 | 0.75${\pm}$0.07 | 0.75${\pm}$0.09 | 0.60${\pm}$0.14

##  4 Discussion & future work

Contrasting image classification and imitation learning datasets. The use of
self-supervised representation learning for pretraining has met with notable
success in image classification [2]. By comparison, results from RL literature
have been mixed, with some positive results, but also several works [4, 5]
which claim that RepL adds little value relative to image augmentation—a
result which we observe in imitation as well. Given this, it’s natural to
wonder _why_ successes from supervised learning have not been reproduced in
sequential decision making problems such as RL and imitation.

Figure 2: We show a sample of STL-10, MAGICAL, and Procgen images. Images on
the same row have the same label (bird, car, etc.) or expert action (up, down,
etc.). It can be easier to tell whether two images have the same label in
classification than in IL tasks.

The case of Behavioural Cloning (BC) is particularly illustrative. BC uses the
same optimization algorithms, loss types, and network architectures as other
forms of image classification, so if RepL is less helpful for BC than for
other forms of classification then it must be due to the choice of training
and evaluation data. For the sake of illustrating differences in data
distributions, Fig. 2 compares the STL-10 dataset—a typical image
classification task—with datasets for MAGICAL and Procgen. dm_control is not
pictured because it has a continuous action space, so there was not a natural
separation of images by action along the $y$ axis.

One notable difference in Fig. 2 is that there is less between-class variation
in MAGICAL and Procgen than in STL-10: the choice of action is often
influenced by fine-grained, local cues in the environment, rather than the
most visually salient axes of variation (background, mean color, etc.). For
example, in MAGICAL the sets of states that correspond to the “forward” and
“left” demonstrator actions cover a similar visual range. Indeed, the agent’s
choice between “left” and “right” could change if its heading shifted by just
a few degrees, even though this visual change would not be obvious to a human.
In contrast, STL-10 exhibits substantial between-class variation: it’s hard to
confuse a the sky-blue background and metal texture of a plane for the natural
setting and fur of a deer. Thus, a RepL method that simply captures the most
visually salient differences between classes may be much more useful for
classification on STL-10 than for control on MAGICAL or Procgen.

Figure 3: t-SNE embedding of representations from a VAE encoder on CoinRun,
labeled with the corresponding actions (left) and discretized returns
(middle). Returns are estimated by applying GAE to an expert PPO demonstrator,
then discretized by rounding to the nearest whole number to produce a “label”.
We compare these with STL-10 image representations generated by a ResNet50
pretrained with SimCLR, colored by class (right).

What is the right downstream prediction target? Both GAIL and BC attempt to
learn a policy that predicts expert actions from observations. We’ve argued
that RepL algorithms may be focusing primarily on the most visually salient
differences between states, at the expense of the fine-grained features
necessary for action prediction. However, it could be that reward- and value-
prediction benefit more from a representation that captures mostly coarse-
grained visual differences. Moreover, Yang and Nachum [26] have observed that
state-based (as opposed to image-based) offline Q-learning does benefit from
existing RepL techniques, even though state-based BC does not. Together, these
facts suggest existing RepL methods might be more helpful when the downstream
prediction target is value or reward rather than action.

To explore this hypothesis, we visualize how well RepL-learned representations
align with action labels, estimated expert returns, and trajectory IDs. In
Figure 3 we show t-SNE projections of observation embeddings taken from seven
expert CoinRun demonstrations. The embeddings were generated by a VAE-
pretrained encoder. We compare these with t-SNE clusters generated from a
ResNet50 with SimCLR on ImageNet, then evaluated on STL-10 (a resized ImageNet
subset).

Representations from a well-trained encoder should cluster nicely according to
the label (e.g. classes, actions) used for the downstream task. We see this
with the STL-10 embeddings, which cluster nicely by class. In contrast, we see
that our encoders for CoinRun do not produce embeddings that cluster nicely by
action. However, they do seem to cluster readily by estimated expert returns.
This is likely a consequence of the events that cause states to have high
value—such as being close to the far wall with the coin—depend primarily on
coarse-grained features of the state. We speculate that this is likely true in
MAGICAL, too, where the reward function tends to depend only on salient
features like whether the agent is overlapping with any of the colored goal
regions.

Our negative results for GAIL and RepL provide reason to be cautious about our
conjecture that reward functions (and value functions) are more amenable to
RepL. A GAIL discriminator is similar to a reward function, but the overall
performance of GAIL does not change much when pretraining the discriminator
with RepL. On the other hand, it is worth noting that the GAIL discriminator
does not in general converge to a valid reward function for the task, so this
is not a direct test of the hypothesis that reward learning is more amenable
to RepL pretraining than policy learning. We therefore believe it is still
worth investigating whether imitation learning algorithms that directly learn
reward functions [27] or value functions [28] benefit more from RepL than
algorithms that learn policies.

The importance of using diverse benchmark tasks. Our experiment results in
Table 5 showed much greater benefit for RepL on DMC than on Procgen and
MAGICAL. This underscores the importance of evaluating across multiple
benchmarks: had we only used DMC, we might have erroneously concluded that
RepL is typically helpful for BC.

The finger-spin (DMC) and CoinRun (Procgen) tasks provide a useful
illustration of how differences in performance across tasks can arise. Fig. 4
shows example saliency maps [29] generated by SimCLR-pretrained encoders in
these two tasks. In finger-spin, the SimCLR encoder mostly attends to
foreground objects, while in CoinRun it attends to the background. This makes
sense: the boundary between the background and terrain is easy to detect and
shifts rapidly as the agent moves, so paying attention to the shape of
background is quite helpful for distinguishing between frames. Unfortunately,
semantically important foreground features in CoinRun, such as obstacles and
gold, are less discriminative, which is why we believe SimCLR is not
dedicating as much model capacity to them. In contrast, the background in
finger-spin changes very little, so SimCLR is forced to attend to foreground
objects that change position between frames.

More generally, we believe that differences between RepL performance across
tasks are due to implicit assumptions that our (unsupervised) RepL algorithms
make about what kinds of features are important. For tasks that do not match
these assumptions, the representation learning algorithms will do poorly,
regardless of how much data is available. In our SimCLR example, information
about background shapes crowds out task-relevant cues like the distance
between the agent and an obstacle. It is therefore important for future
research to (1) consider whether the implicit assumptions underlying a given
RepL algorithm are likely to help models acquire useful invariances for the
desired tasks; and (2) test on multiple domains to ensure that the claimed
improvements are robust across environments.

Figure 4: Saliency map generated by an encoder trained using SimCLR. Top row
shows input frames, averaged across a three-frame stack of inputs. Bottom row
shows saliency map overlayed on top of grayscale images, with darker blue
shading indicating greater influence over the network’s output. Notice that
SimCLR attends mainly to the foreground in DMC, and mainly to the background
in CoinRun.

##  5 Conclusion

We have seen that, when compared against a well-tuned IL baseline using image
augmentations, the impacts of representation learning for imitation are
limited. On some benchmark suites it appears that it helps, while on others
there is not much impact, suggesting that the effect of RepL is quite
benchmark-specific. Our analysis has identified several hypotheses that could
help understand _when_ and _where_ representation learning can be useful. We
are excited to see future work investigate these hypotheses, and hope the
EIRLI framework can serve as a useful starting point for any such
investigations.

## Acknowledgments and Disclosure of Funding

The authors would like to thank Michael Chang, Ben Eysenbach, Aravind
Srinivas, and Olivia Watkins for feedback on earlier versions of this work.
This work was supported in part by the DOE CSGF under grant number DE-
SC0020347, along with a grant from the Open Philanthropy Project and
computational support from Google.

## References

  * Jeon et al. [2020] Hong Jun Jeon, Smitha Milli, and Anca D Dragan.  Reward-rational (implicit) choice: A unifying formalism for reward learning.  _arXiv preprint arXiv:2002.04833_ , 2020.
  * Chen et al. [2020a] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  _Proceedings of the 37th International Conference on Machine Learning_ , 2020a.
  * Lee et al. [2020] Kuang-Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama.  Predictive information accelerates learning in rl.  _arXiv preprint arXiv:2007.12401_ , 2020.
  * Laskin et al. [2020a] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas.  Reinforcement learning with augmented data.  _arXiv preprint arXiv:2004.14990_ , 2020a.
  * Kostrikov et al. [2020] Ilya Kostrikov, Denis Yarats, and Rob Fergus.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  _arXiv preprint arXiv:2004.13649_ , 2020.
  * Toyer et al. [2020] Sam Toyer, Rohin Shah, Andrew Critch, and Stuart Russell.  The MAGICAL benchmark for robust imitation.  _arXiv preprint arXiv:2011.00401_ , 2020.
  * Cobbe et al. [2020] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman.  Leveraging procedural generation to benchmark reinforcement learning.  In _International conference on machine learning_ , pages 2048–2056. PMLR, 2020.
  * Tassa et al. [2018] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al.  DeepMind control suite.  _arXiv preprint arXiv:1801.00690_ , 2018.
  * Kingma and Welling [2013] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Hendrycks et al. [2019] Dan Hendrycks, Norman Mu, Ekin Dogus Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshminarayanan.  Augmix: A simple data processing method to improve robustness and uncertainty.  In _International Conference on Learning Representations_ , 2019.
  * Sohn et al. [2020] Kihyuk Sohn, David Berthelot, Chun-Liang Li, Zizhao Zhang, Nicholas Carlini, Ekin D Cubuk, Alex Kurakin, Han Zhang, and Colin Raffel.  FixMatch: Simplifying semi-supervised learning with consistency and confidence.  _In Proceedings of the Thirty-fourth Conference on Neural Information Processing Systems_ , 2020.
  * van den Oord et al. [2018] Aaron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_ , 2018.
  * He et al. [2020] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9729–9738, 2020.
  * Chen et al. [2020b] Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton.  Big self-supervised models are strong semi-supervised learners.  _arXiv preprint arXiv:2006.10029_ , 2020b.
  * Grill et al. [2020] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent: A new approach to self-supervised learning.  _arXiv preprint arXiv:2006.07733_ , 2020.
  * Ha and Schmidhuber [2018] David Ha and Jürgen Schmidhuber.  World models.  _arXiv preprint arXiv:1803.10122_ , 2018.
  * Hafner et al. [2019] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson.  Learning latent dynamics for planning from pixels.  In _International Conference on Machine Learning_ , pages 2555–2565. PMLR, 2019.
  * Laskin et al. [2020b] Michael Laskin, Aravind Srinivas, and Pieter Abbeel.  Curl: Contrastive unsupervised representations for reinforcement learning.  _Proceedings of the 37th International Conference on Machine Learning, Vienna, Austria, PMLR 119_ , 2020b.  arXiv:2003.06417.
  * Stooke et al. [2020] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin.  Decoupling representation learning from reinforcement learning.  _arXiv preprint arXiv:2009.08319_ , 2020.
  * Schmidhuber [1991] Jürgen Schmidhuber.  Reinforcement learning in markovian and non-markovian environments.  In _Advances in Neural Information Processing Systems 3, NIPS’3_ , pages 500–506, 1991.
  * Poole et al. [2019] Ben Poole, Sherjil Ozair, Aaron van den Oord, Alexander A Alemi, and George Tucker.  On variational bounds of mutual information.  _Proceedings of the 36th International Conference on Machine Learning_ , 2019.
  * Fischer [2020] Ian Fischer.  The conditional entropy bottleneck.  _Entropy_ , 2020.
  * Pomerleau [1991] Dean A Pomerleau.  Efficient training of artificial neural networks for autonomous navigation.  _Neural Computation_ , 3(1):88–97, 1991.
  * Ho and Ermon [2016] Jonathan Ho and Stefano Ermon.  Generative adversarial imitation learning.  In _Advances in neural information processing systems_ , pages 4565–4573, 2016.
  * Zolna et al. [2019] Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang.  Task-relevant adversarial imitation learning.  _arXiv preprint arXiv:1910.01077_ , 2019.
  * Yang and Nachum [2021] Mengjiao Yang and Ofir Nachum.  Representation matters: Offline pretraining for sequential decision making.  _arXiv preprint arXiv:2102.05815_ , 2021.
  * Fu et al. [2017] Justin Fu, Katie Luo, and Sergey Levine.  Learning robust rewards with adversarial inverse reinforcement learning.  _arXiv preprint arXiv:1710.11248_ , 2017.
  * Reddy et al. [2019] Siddharth Reddy, Anca D Dragan, and Sergey Levine.  SQIL: Imitation learning via reinforcement learning with sparse rewards.  _arXiv preprint arXiv:1905.11108_ , 2019.
  * Simonyan et al. [2013] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman.  Deep inside convolutional networks: Visualising image classification models and saliency maps.  _arXiv preprint arXiv:1312.6034_ , 2013.
  * Igl et al. [2018] Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson.  Deep variational reinforcement learning for pomdps.  _arXiv preprint arXiv:1806.02426_ , 2018.
  * Lee et al. [2019] Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine.  Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model.  _arXiv preprint arXiv:1907.00953_ , 2019.
  * Agrawal et al. [2016] Pulkit Agrawal, Ashvin Nair, Pieter Abbeel, Jitendra Malik, and Sergey Levine.  Learning to poke by poking: Experiential learning of intuitive physics.  _arXiv preprint arXiv:1606.07419_ , 2016.
  * Yingjun and Xinwen [2019] Pei Yingjun and Hou Xinwen.  Learning representations in reinforcement learning: An information bottleneck approach.  _arXiv preprint arXiv:1911.05695_ , 2019.
  * Mazoure et al. [2020] Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R Devon Hjelm.  Deep reinforcement and infomax learning.  _arXiv preprint arXiv:2006.07217_ , 2020.
  * Guo et al. [2020] Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean-bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar.  Bootstrap latent-predictive representations for multitask reinforcement learning.  _arXiv preprint arXiv:2004.14646_ , 2020.
  * Schwarzer et al. [2020] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman.  Data-efficient reinforcement learning with momentum predictive representations.  _arXiv preprint arXiv:2007.05929_ , 2020.
  * Zhu et al. [2020] Jinhua Zhu, Yingce Xia, Lijun Wu, Jiajun Deng, Wengang Zhou, Tao Qin, and Houqiang Li.  Masked contrastive representation learning for reinforcement learning.  _arXiv preprint arXiv:2010.07470_ , 2020.
  * Yu et al. [2021] Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, and Zhibo Chen.  PlayVirtual: Augmenting cycle-consistent virtual trajectories for reinforcement learning.  _arXiv preprint arXiv:2106.04152_ , 2021.
  * He et al. [2015] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 1026–1034, 2015.
  * Fischer and Alemi [2020] Ian Fischer and Alex Alemi.  Ceb improves model robustness.  _Entropy_ , 2020.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes]

    2. (b)

Did you describe the limitations of your work? [Yes]

    3. (c)

Did you discuss any potential negative societal impacts of your work? [Yes]
Both the limitations and potential societal impacts are discussed in the
appendix.

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [N/A]

    2. (b)

Did you include complete proofs of all theoretical results? [N/A]

  3. 3.

If you ran experiments (e.g. for benchmarks)…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes]

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [Yes]

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]  The
information is provided in the Appendix.

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes]

    2. (b)

Did you mention the license of the assets? [Yes]  We use the MIT licence for
our codebase, discussed in the Appendix.

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [Yes]

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [N/A]

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [N/A]

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]

##  Appendix A Design choices for reinforcement learning

Table A1: Design choices made in representation learning for reinforcement learning. Act, Aug, Mom, Proj and Comp respectively show whether action conditioning, augmentation, momentum, projection heads, and compression were used. P/J determines whether the representation learning is an initial (P)retraining step, or is (J)ointly learned alongside reinforcement learning. R/C/B/N in the Task column refer to Reconstruction, Contrastive, Bootstrap, or None. Note that different papers may use different sets of augmentations.  Algorithm | Task | RL alg. | Context | Target | Act | Aug | Mom | Proj | Comp | P/J
---|---|---|---|---|---|---|---|---|---|---
World models [16] | R | CMA-ES | $o_{t}$ | $o_{t},o_{t+1}$ | ✓ | ✗ | ✗ | ✗ | ✗ | P
DVRL [30] | R | A2C | $o_{t}$ | $o_{t+k}$ | ✓ | ✗ | ✗ | ✗ | ✗ | J
PlaNet [17] | R | MPC + CEM | $o_{1:t}$ | $o_{t+1:T},r_{t+1:T}$ | ✓ | ✗ | ✗ | ✗ | ✗ | J
SLAC [31] | R | SAC | $o_{t}$ | $o_{t+1}$ | ✓ | ✗ | ✗ | ✗ | ✗ | J
Poke [32] | R | - | $o_{t},o_{t+1}$ | $a_{t}$ | - | ✗ | ✗ | ✗ | ✗ | -
RAD [4] | N | PPO, SAC | $o_{t}$ | - | - | ✓ | - | - | ✗ | J
DrQ [5] | N | DQN, SAC | $o_{t}$ | - | - | ✓ | - | - | ✗ | J
CURL [18] | C | DQN, SAC | $o_{t}$ | $o_{t}$ | ✗ | ✓ | ✓ | ✗ | ✗ | J
CPC [12] | C | A2C | $o_{t}$ | $o_{t+k}$ | ✗ | ✗ | ✗ | ✗ | ✗ | J
Bottleneck [33] | C | A2C | $o_{t}$ | $o_{t+k}$ | ✓ | ✗ | ✗ | ✗ | ✗ | J
DRIML [34] | C | C51 | $o_{t}$ | $o_{t+k}$ | ✓ | ✗ | ✗ | ✗ | ✗ | J
PI-SAC [3] | C | SAC | $o_{t}$ | $o_{t+k},r_{t+k}$ | ✓ | ✓ | ✓ | ✓ | ✓ | J
ATC [19] | C | SAC, PPO | $o_{t}$ | $o_{t+k}$ | ✗ | ✓ | ✓ | ✓ | ✗ | P
PBL [35] | B | PopArt-IMPALA | $o_{1:t}$ | $o_{t+k}$ | ✓ | ✗ | ✗ | ✓ | ✗ | J
SPR [36] | B | DQN | $o_{t}$ | $o_{t+1:T}$ | ✓ | ✓ | ✓ | ✓ | ✗ | J
M-CURL [37] | C | DQN, SAC | $o_{1:t}$ | $o_{t}$ | ✗ | ✓ | ✓ | ✗ | ✗ | J
PlayVirtual [38] | B | DQN, SAC | $o_{t}$ | $o_{t+1:T}$ | ✓ | ✓ | ✓ | ✓ | ✗ | J

##  Appendix B Additional information on experiment setup

###  B.1 Environment setup

DMC. The training set for each DMC task consist of $250$ trajectories produced
by an expert policy trained with RAD [4]. For each of our methods we report
the mean return for the final policy, which ranges between 0 and 1,000.

Procgen. For each Procgen task, we reserve 100 environment seeds as “training
levels”, then use a separate 100 seeds as “testing levels” to evaluate
generalization. The training set for each task consists of 110K frames
produced by applying RAD’s trained agent to the training levels. We report
mean returns on the train and test levels separately in our results.

MAGICAL. We give the IL and RepL algorithms access to five human
demonstrations from the demo variant, then provide RepL with an additional
150,000 time steps of random rollouts (1,250 to 3,750 trajectories, depending
on the environment). This simulates situations where human expert
demonstrations are expensive to collect, but robot exploration is relatively
cheap. For space reasons, we report mean scores averaged across all variants,
which range between 0 and 1; full results are in Appendix G.

###  B.2 RepL algorithms

As is discussed in the main text, we evaluated 5 RepL algorithms. SimCLR and
TemporalCPC are contrastive baselines: SimCLR must assign similar
representations to augmented copies of the same observation, and different
representations to augmented copies of different observations. TemporalCPC
must additionally account for dynamics by assigning similar representations to
any pair of frames that are separated by a gap of $\Delta t$ time steps (we
use $\Delta t=8$). Like SimCLR, the VAE attempts to represent one frame at a
time, without using temporal offsets. We also include methods that explicitly
condition on or generate actions: dynamics predicts the next observation given
the current observation and action, while inverse dynamics predicts which
action was used to transition between two adjacent states. Note that we do not
use momentum, projection heads, or compression in our final experiments, since
our preliminary experiments did not show a significant advantage to doing so.

##  Appendix C Related benchmarking work in RL and IL

Table A2: Summary of previous works comparing different representation learning (repL) algorithms on imitation learning (IL) and reinforcement learning (RL). These works differ along a few axis on whether they provide a design breakdown of repL algorithms, their area of focus, the number of benchmarks covered, the number of algorithms experimented, whether they evaluate image-based environments, and whether they compare their results with an image augmentation baseline.  Research | Design Breakdown | Area | # of Benchmarks | # of Algo. | Image env. | Aug
---|---|---|---|---|---|---
ACL [26] | ✓ | RL, IL | 1 | 10 | ✗ | ✗
DrQ [5] | ✗ | RL | 1 | 6 | ✓ | ✓
RAD [4] | ✗ | RL | 3 | 8 | ✓ | ✓
EIRLI (ours) | ✓ | IL | 3 | 7 | ✓ | ✓

To the best of our knowledge, there are three existing work that survey RepL
methods on imitation learning and reinforcement learning. We summarize their
differences in Table A2. ACL [26] provided a design breakdown of different
RepL algorithms too, and they found that many RepL algorithms perform poorly
on imitation learning but can provide extra benefits for offline RL. DrQ [5]
and RAD [4] both discussed the effect of image augmentation in reinforcement
learning in great detail, and showed that with well-tuned image augmentations,
a standard reinforcement learning framework can outperform many self-
supervised learning methods on RL.

##  Appendix D Hyperparameter details

Hyperparameter | Value | Tuning range
---|---|---
All algorithms
Optimizer | Adam | -
LR | $10^{-4}$ |  $10^{-6}$–$10^{-2}$
Training batches | 5,000 | -
Representation dim. | 128 | 64–256
VAE, dyn., inv. dyn.
Batch size | 64 | -
Augmentations | - | -
VAE
VAE $\beta$ | $10^{-6}$ |  $10^{-7}$–$1.0$
SimCLR, TCPC
Batch size | 384 | 64–512
Augmentations | trans. rot., blur, col. jit. | -
TCPC
Temporal offset | 8 steps | -
Table A3: Hyperparameters for representation learning. Note that for joint training, BC and RepL use the same optimizer, and thus have the same learning rate. Hyperparameter | Value | Tuning range
---|---|---
All benchmarks
Optimizer | Adam | -
LR | $10^{-4}$ | -
Entropy coeff. | $10^{-3}$ | -
$\ell_{2}$ reg. coeff. | $10^{-5}$ | -
Augmentations | trans., rot., blur, col. jit. | -
All benchmarks, pretraining
Batch size | 32 | -
All benchmarks, joint training
Batch size | 64 | -
dm_control and Procgen, pretraining and joint training
Training batches | 1M | 1M–4M
MAGICAL, pretraining
Training batches | 20k | 5k–20k
MAGICAL, joint training
Training batches | 30k | -
Table A4:  Hyperparameters for behavioral cloning. Sections marked “pretraining” show the hyperparameters used for BC after pretraining; sections marked “joint training” apply to BC during joint training. Representation learning hyperparameters, such as the batch size, are covered separately in Table A3. Hyperparameter | DMC | Procgen | MAGICAL | Tuning range
---|---|---|---|---
Policy (PPO) |  |  |  |
# parallel envs | 32 | 32 | 32 | -
Time steps per round | 8 | 10 | 7 | 4–12
Epochs per round | 12 | 9 | 7 | 4–12
Adam minibatch size | 48 | 48 | 48 | -
Initial Adam step size | $10^{-4}$ | $10^{-4}$ | $2.5\times 10^{-4}$ |  $5\times(10^{-5}$–$10^{-4})$
Final Adam step size | Linearly annealed to 0 over training | -
Discount $\gamma$ | 0.99 | 0.6 | 0.99 | 0.6–1
GAE $\lambda$ | 0.8 | 0.6 | 0.76 | 0.6–0.9
Entropy bonus | $10^{-8}$ | $5\times 10^{-6}$ | $4.5\times 10^{-8}$ |  $10^{-10}$–$10^{-3}$
Advantage clip $\epsilon$ | 0.02 | 0.01 | 0.006 |  $0.001$–$0.1$
Grad. clip $\ell_{2}$ norm | 1 | 1 | 1 | -
Augmentations | - | - | - | -
Discriminator |  |  |  |
Batch size | 48 | 48 | 48 | -
Adam step size | $10^{-3}$ | $2.5\times 10^{-3}$ | $5.7\times 10^{-4}$ |  $5\times(10^{-4}$–$10^{-3})$
Disc. steps per round | 6 | 2 | 2 | 1–8
Augmentations |  Erase, blur, noise, rot. |  Col. jit., flip LR, blur, noise, rot., trans. |  Col. jit., erase, flip LR, blur, noise, rot. |  Col. jit., erase, flip LR, blur, noise, rot., trans.
Misc. |  |  |  |
Total env. steps of training | $5\times 10^{5}$ | $5\times 10^{5}$ | $5\times 10^{5}$ | -
Reward norm. std. dev. | 0.01 | 0.01 | 0.01 | -
Table A5:  Hyperparameters for GAIL experiments. We use the word “round” to
describe the repeated sequence of data collection, followed by PPO updates on
the collected data, followed by discriminator updates on both demonstrations
and rollouts. Representation learning hyperparameters, such as the batch size,
are covered separately in Table A3.

Environments and datasets For each dm_control environment, we generated
synthetic demonstration data using RAD with default algorithm hyperparameters
[4].111https://github.com/MishaLaskin/rad Environment configurations (such as
action repeat, frame stack, etc.) were the same for both RAD and our IL
algorithms. Specifically:

  * •

In cheetah-run, we used an action repeat of 4, resulting in a trajectory
length of $1000/4=250$. Our demonstration dataset consisted of 250
trajectories (62,500 time steps) from the RAD demonstration agent, with a mean
return of $\approx$827 (recall that return ranges between 0 and 1,000 for all
DMC environments).

  * •

In finger-spin, we used an action repeat of 2, resulting in a trajectory
length of $1000/2=500$. Our dataset again consisted of 250 trajectories
(125,000 time steps) sampled from the RAD demonstration agent, with mean
return of $\approx$963.

  * •

In reacher-easy, we used an action repeat of 8, resulting in a trajectory
length of $1000/8=125$. Our dataset of 250 trajectories (31,250 time steps)
had mean return $\approx$977, and was again generated by RAD.

For all DMC environments, we used a frame stack of 3.

As with DMC, we generated expert demonstrations for Procgen using a policy
trained with RAD.222https://github.com/pokaxpoka/rad_procgen We used the easy
variants of all environments, with a frame stack of 3 and no action repeat. We
used a demonstration dataset of around 114,000 timesteps for each agent. The
mean trajectory lengths and returns are as follows:

  * •

For CoinRun, trajectories averaged 26 steps, and the demonstrator had an
average return of 8.7.

  * •

For Fruitbot, trajectories had an average length of 442, and the demonstrator
attained a mean return of 29.75.

  * •

For Jumper, trajectories had an average length of 76, and mean return of 8.7.

For each MAGICAL environment, we used a fixed subset of five demonstration
trajectories (initially selected at random) from the human dataset provided
with the benchmark [6]. We used egocentric views with a frame stack of four
and no action repeat. Because there was no action repeat, trajectory lengths
remained at the defaults for the benchmark suite: 40 for MoveToRegion, 80 for
MoveToCorner, 120 for MatchRegions. For each benchmark, we used between 25 and
28 demonstration trajectories, and the demonstrator attained the maximum
return of 1.0 (on a 0.0–1.0 scale) in each trajectory. In addition to
demonstrations, our MAGICAL experiments also used random rollout datasets of
150,000 timesteps, all generated by uniformly sampling from the action set at
each time step. This equates to between 1,250 and 3,750 trajectories,
depending on the horizon of the task.

RepL hyperparameters Representation learning hyperparameters are given in
Table A3. Note that the contrastive algorithms have slightly different
hyperparameters from the other RepL algorithms. We found that a batch size
close to 400 was important for contrastive algorithm performance; setting this
value too low or too high (e.g. 500+) decreased performance. Predictive and
reconstructive algorithms are less sensitive to batch size, so we used a
computationally convenient batch size (64).

For the contrastive algorithms, we used a mixture of translation (trans.),
rotation (rot.), Gaussian blur (blur), and color jitter (col. jit.)
augmentations. The translation augmentation translates the image by up to 5%
of image dimensions; the rotation augmentation rotates the image by up to 5∘;
the Gaussian blur augmentation applies a Gaussian blur kernel with
$\sigma=1$px; and the color jitter augmentation randomizes the hue by up to
0.15 radians. We did not find the algorithms were highly sensitive to the
choice of augmentations, but these augmentations did appear to perform
fractionally better than the other choices that we considered during manual
tuning. For non-contrastive algorithms, we did not use augmentations.

For the VAE, we used a mean squared error loss for reconstruction, and down-
weighted the KL prior term by a factor of $\beta$. Specifically, our loss was

$\mathcal{L}_{\text{VAE}}=\frac{1}{n}\sum_{i=1}^{n}(x_{i}-y_{i})^{2}-\beta\operatorname{KL}(e_{\theta}(\cdot\mid
x_{i})\|\mathcal{N}(\cdot;0,I))~{},$

where $i$ indexes over the $n$ elements of the output image. This is
equivalent to using a traditional $\beta$-VAE with a fixed variance of
$\sigma^{2}=\frac{n}{2}$ for the Gaussian output distribution.

IL hyperparameters Imitation learning hyperparameters are given in Table A4
for BC, and Table A5 for GAIL. IL hyperparameters were the same for both
control and RepL runs, except for the network initialization, where RepL runs
initialized from the RepL-trained encoder, while control runs used a fresh He
initialization [39]. We found that DMC and Procgen needed substantially more
training than MAGICAL; we found that MAGICAL performance was plateauing even
with an order of magnitude less training.

Compute information Running one experiment (a single seed of combined RepL and
IL) on DMC and Procgen for 1M batch updates takes about 40 hours on one NVIDIA
1080Ti, and running one for MAGICAL between 30 minutes (for 20,000-batch
control without augmentations) and 10 hours (for a contrastive method using
joint training for 30,000 batches). Generating the results in this paper takes
around 6,600 GPU hours for DMC, 8,800 GPU hours for Procgen, and 26–525 hours
for MAGICAL (assuming 4 seeds per GPU, and 30 minutes to 10 hours per seed).

##  Appendix E Limitations, social impacts, and benchmark license

##### Limitations

The main limitation of our findings is that we only investigate policy
learning with BC. Our findings therefore may not generalise to IL algorithms
that learn more than just a policy. This includes IRL algorithms, which
typically learn both a reward function and a policy, as well as IL algorithms
like SQIL [28] that learn a Q function rather than directly learning a policy.

##### Social impacts

We do not foresee any negative near-term social impact from our work.

##### License

We release our codebase and associated data under the MIT license.

##  Appendix F Implementation of components in the codebase

In Section 2 we analyzed several design axes and their components. We
elaborate in this section our current implementation status of these
components in the codebase.

Target selection. Different versions of this design choice are implemented by
inheriting from the TargetPairConstructor class within the codebase. Currently
implemented are IdentityPairConstructor, in which the context and target are
identical, and TemporalOffsetPairConstructor, which can be given a desired
temporal offset, and configured to optionally return the action as extra
context.

Loss type. Different versions of loss functions are implemented by inheriting
from RepresentationLoss. We have already implemented a wide variety of losses,
including VAE, mean squared error, negative log likelihood, CEB, and several
contrastive losses.

Augmentation. This design choice is implemented in subclasses of Augmenter. We
rely on a standard library to implement the augmentations. Each subclass
augments a different set of inputs: both the context and target, only the
context, or neither the context nor target.

Encoder. Different versions of encoder are implemented by inheriting from the
Encoder class within the codebase. We have implemented encoders that work on
individual images, as well as a RecurrentEncoder.

Decoder. Different versions of decoder are implemented by inheriting from the
LossDecoder class within the codebase. Currently implemented decoders support
image reconstruction, action conditioning, and projection heads.

##  Appendix G Complete MAGICAL results

Task Dynamics InvDyn SimCLR TemporalCPC VAE BC, augs BC, no augs MatchRegions-
Demo 0.71${\pm}$0.06 0.72${\pm}$0.07 0.72${\pm}$0.06 0.68${\pm}$0.03
0.70${\pm}$0.05 0.76${\pm}$0.05 0.63${\pm}$0.15 -TestDynamics 0.55${\pm}$0.09
0.56${\pm}$0.07 0.56${\pm}$0.07 0.52${\pm}$0.01 0.56${\pm}$0.08
0.66${\pm}$0.05 0.38${\pm}$0.12 -TestColour 0.71${\pm}$0.06* 0.70${\pm}$0.08*
0.67${\pm}$0.04* 0.66${\pm}$0.04* 0.67${\pm}$0.06* 0.53${\pm}$0.04
0.30${\pm}$0.08 -TestShape 0.61${\pm}$0.06 0.64${\pm}$0.09 0.62${\pm}$0.06
0.61${\pm}$0.04 0.60${\pm}$0.05 0.70${\pm}$0.03 0.43${\pm}$0.13 -TestJitter
0.70${\pm}$0.06 0.65${\pm}$0.06 0.66${\pm}$0.07 0.65${\pm}$0.04
0.68${\pm}$0.03 0.68${\pm}$0.05 0.41${\pm}$0.13 -TestLayout 0.04${\pm}$0.01
0.04${\pm}$0.01 0.04${\pm}$0.01 0.05${\pm}$0.01 0.04${\pm}$0.01
0.05${\pm}$0.01 0.05${\pm}$0.03 -TestCountPlus 0.04${\pm}$0.02 0.04${\pm}$0.01
0.06${\pm}$0.03 0.05${\pm}$0.01 0.04${\pm}$0.02 0.05${\pm}$0.02
0.04${\pm}$0.02 -TestAll 0.04${\pm}$0.02 0.05${\pm}$0.02 0.05${\pm}$0.03
0.06${\pm}$0.01* 0.04${\pm}$0.02 0.05${\pm}$0.01 0.03${\pm}$0.02 Average
0.42${\pm}$0.04 0.42${\pm}$0.04 0.42${\pm}$0.03 0.41${\pm}$0.01
0.42${\pm}$0.03 0.43${\pm}$0.02 0.28${\pm}$0.08 MoveToCorner-Demo
0.94${\pm}$0.08 0.92${\pm}$0.07 0.88${\pm}$0.09 0.86${\pm}$0.10
0.89${\pm}$0.08 0.86${\pm}$0.08 0.99${\pm}$0.01* -TestDynamics 0.83${\pm}$0.05
0.87${\pm}$0.05* 0.85${\pm}$0.04* 0.80${\pm}$0.06 0.75${\pm}$0.04
0.76${\pm}$0.08 0.68${\pm}$0.04 -TestColour 0.90${\pm}$0.10 0.87${\pm}$0.06
0.86${\pm}$0.13 0.86${\pm}$0.06 0.87${\pm}$0.07 0.75${\pm}$0.16
0.85${\pm}$0.07 -TestShape 0.90${\pm}$0.06 0.88${\pm}$0.07 0.91${\pm}$0.05*
0.87${\pm}$0.02 0.84${\pm}$0.09 0.84${\pm}$0.06 0.71${\pm}$0.10 -TestJitter
0.81${\pm}$0.09 0.80${\pm}$0.04 0.84${\pm}$0.04* 0.79${\pm}$0.04
0.77${\pm}$0.04 0.78${\pm}$0.00 0.58${\pm}$0.08 -TestAll 0.66${\pm}$0.13
0.64${\pm}$0.09 0.65${\pm}$0.07 0.62${\pm}$0.06 0.57${\pm}$0.12
0.67${\pm}$0.09 0.51${\pm}$0.13 Average 0.84${\pm}$0.07 0.83${\pm}$0.04
0.83${\pm}$0.04* 0.80${\pm}$0.02 0.78${\pm}$0.06 0.78${\pm}$0.05
0.72${\pm}$0.04 MoveToRegion-Demo 0.99${\pm}$0.01 0.99${\pm}$0.01
1.00${\pm}$0.01 0.98${\pm}$0.02 0.99${\pm}$0.02 1.00${\pm}$0.00
1.00${\pm}$0.00 -TestDynamics 0.99${\pm}$0.01 1.00${\pm}$0.01 0.99${\pm}$0.01
0.99${\pm}$0.01 0.99${\pm}$0.01 1.00${\pm}$0.01 1.00${\pm}$0.00 -TestColour
0.99${\pm}$0.02* 0.99${\pm}$0.01* 1.00${\pm}$0.00* 0.99${\pm}$0.01*
0.97${\pm}$0.04* 0.54${\pm}$0.08 0.84${\pm}$0.13* -TestJitter 0.99${\pm}$0.01
1.00${\pm}$0.00 1.00${\pm}$0.00 1.00${\pm}$0.00 1.00${\pm}$0.01
1.00${\pm}$0.00 0.99${\pm}$0.01 -TestLayout 0.52${\pm}$0.05 0.51${\pm}$0.05
0.45${\pm}$0.03 0.44${\pm}$0.07 0.46${\pm}$0.11 0.64${\pm}$0.03
0.64${\pm}$0.11 -TestAll 0.45${\pm}$0.06* 0.50${\pm}$0.09* 0.46${\pm}$0.06*
0.42${\pm}$0.08* 0.47${\pm}$0.12* 0.28${\pm}$0.06 0.38${\pm}$0.08* Average
0.82${\pm}$0.02* 0.83${\pm}$0.02* 0.82${\pm}$0.01* 0.81${\pm}$0.01*
0.81${\pm}$0.05* 0.74${\pm}$0.02 0.81${\pm}$0.04*

Table A6:  Complete results for all variants of the evaluated MAGICAL tasks,
with BC plus RepL pretraining. Refer to Table A7 for joint training.

Task Dynamics InvDyn SimCLR TemporalCPC VAE BC, augs BC, no augs MatchRegions-
Demo 0.78${\pm}$0.03 0.51${\pm}$0.16 0.73${\pm}$0.06 0.01${\pm}$0.01
0.72${\pm}$0.09 0.78${\pm}$0.06 0.70${\pm}$0.04 -TestDynamics 0.62${\pm}$0.04
0.39${\pm}$0.14 0.59${\pm}$0.04 0.01${\pm}$0.01 0.58${\pm}$0.05
0.63${\pm}$0.05 0.40${\pm}$0.03 -TestColour 0.60${\pm}$0.03* 0.17${\pm}$0.07
0.55${\pm}$0.06* 0.01${\pm}$0.01 0.48${\pm}$0.04 0.47${\pm}$0.08
0.35${\pm}$0.10 -TestShape 0.71${\pm}$0.04* 0.36${\pm}$0.14 0.67${\pm}$0.05
0.01${\pm}$0.01 0.68${\pm}$0.05 0.66${\pm}$0.04 0.48${\pm}$0.02 -TestJitter
0.69${\pm}$0.05 0.36${\pm}$0.14 0.66${\pm}$0.04 0.01${\pm}$0.01
0.65${\pm}$0.04 0.66${\pm}$0.06 0.40${\pm}$0.10 -TestLayout 0.05${\pm}$0.01
0.04${\pm}$0.01 0.04${\pm}$0.01 0.01${\pm}$0.01 0.05${\pm}$0.02
0.07${\pm}$0.01 0.03${\pm}$0.01 -TestCountPlus 0.04${\pm}$0.01 0.02${\pm}$0.01
0.04${\pm}$0.01 0.01${\pm}$0.01 0.05${\pm}$0.03 0.07${\pm}$0.02
0.04${\pm}$0.01 -TestAll 0.04${\pm}$0.01 0.02${\pm}$0.01 0.03${\pm}$0.01
0.01${\pm}$0.01 0.04${\pm}$0.01 0.07${\pm}$0.02 0.04${\pm}$0.03 Average
0.44${\pm}$0.02 0.23${\pm}$0.08 0.41${\pm}$0.02 0.01${\pm}$0.01
0.41${\pm}$0.03 0.43${\pm}$0.03 0.31${\pm}$0.02 MoveToCorner-Demo
0.86${\pm}$0.04 0.37${\pm}$0.25 0.90${\pm}$0.12 0.03${\pm}$0.03
0.94${\pm}$0.05 0.95${\pm}$0.04 0.99${\pm}$0.00* -TestDynamics 0.82${\pm}$0.07
0.35${\pm}$0.27 0.86${\pm}$0.05 0.03${\pm}$0.03 0.87${\pm}$0.04
0.82${\pm}$0.06 0.75${\pm}$0.09 -TestColour 0.84${\pm}$0.08 0.25${\pm}$0.29
0.61${\pm}$0.17 0.01${\pm}$0.01 0.73${\pm}$0.15 0.77${\pm}$0.15
0.76${\pm}$0.17 -TestShape 0.84${\pm}$0.12 0.44${\pm}$0.27 0.88${\pm}$0.09
0.02${\pm}$0.03 0.93${\pm}$0.05 0.89${\pm}$0.04 0.76${\pm}$0.09 -TestJitter
0.76${\pm}$0.06 0.25${\pm}$0.18 0.78${\pm}$0.06 0.02${\pm}$0.03
0.83${\pm}$0.07 0.81${\pm}$0.08 0.55${\pm}$0.12 -TestAll 0.57${\pm}$0.11
0.13${\pm}$0.15 0.53${\pm}$0.10 0.01${\pm}$0.01 0.63${\pm}$0.09
0.57${\pm}$0.10 0.39${\pm}$0.21 Average 0.78${\pm}$0.07 0.30${\pm}$0.22
0.76${\pm}$0.05 0.02${\pm}$0.02 0.82${\pm}$0.06 0.80${\pm}$0.05
0.70${\pm}$0.09 MoveToRegion-Demo 1.00${\pm}$0.00 0.60${\pm}$0.34
1.00${\pm}$0.00 0.81${\pm}$0.14 1.00${\pm}$0.00 1.00${\pm}$0.00
1.00${\pm}$0.00 -TestDynamics 1.00${\pm}$0.00 0.55${\pm}$0.37 1.00${\pm}$0.00
0.81${\pm}$0.14 1.00${\pm}$0.00 1.00${\pm}$0.00 0.99${\pm}$0.01 -TestColour
0.64${\pm}$0.15 0.22${\pm}$0.15 0.93${\pm}$0.07* 0.25${\pm}$0.06
0.69${\pm}$0.12 0.65${\pm}$0.13 0.77${\pm}$0.12 -TestJitter 1.00${\pm}$0.00
0.49${\pm}$0.40 1.00${\pm}$0.00 0.78${\pm}$0.08 1.00${\pm}$0.00
1.00${\pm}$0.00 0.98${\pm}$0.01 -TestLayout 0.64${\pm}$0.05 0.17${\pm}$0.16
0.30${\pm}$0.06 0.10${\pm}$0.05 0.68${\pm}$0.06* 0.61${\pm}$0.04
0.63${\pm}$0.06 -TestAll 0.29${\pm}$0.06 0.10${\pm}$0.08 0.23${\pm}$0.04
0.06${\pm}$0.03 0.28${\pm}$0.03* 0.23${\pm}$0.03 0.31${\pm}$0.08 Average
0.76${\pm}$0.02 0.35${\pm}$0.24 0.74${\pm}$0.01 0.47${\pm}$0.07
0.77${\pm}$0.02 0.75${\pm}$0.02 0.78${\pm}$0.04

Table A7:  Complete results for all variants of the evaluated MAGICAL tasks,
with BC plus RepL joint training. Refer to Table A6 for BC plus RepL
pretraining.

Task Dynamics InvDyn SimCLR TemporalCPC VAE GAIL augs GAIL no augs
MatchRegions-Demo 0.76${\pm}$0.19 0.59${\pm}$0.24 0.83${\pm}$0.09
0.68${\pm}$0.22 0.51${\pm}$0.28 0.81${\pm}$0.10 0.36${\pm}$0.22 -TestDynamics
0.74${\pm}$0.17 0.55${\pm}$0.22 0.79${\pm}$0.11 0.64${\pm}$0.25
0.51${\pm}$0.27 0.77${\pm}$0.13 0.37${\pm}$0.22 -TestColour 0.21${\pm}$0.04
0.18${\pm}$0.06 0.25${\pm}$0.04 0.21${\pm}$0.07 0.14${\pm}$0.08
0.24${\pm}$0.03 0.08${\pm}$0.05 -TestShape 0.72${\pm}$0.20 0.56${\pm}$0.22
0.81${\pm}$0.10 0.61${\pm}$0.25 0.49${\pm}$0.26 0.76${\pm}$0.10
0.35${\pm}$0.20 -TestJitter 0.73${\pm}$0.17 0.57${\pm}$0.20 0.83${\pm}$0.10
0.67${\pm}$0.22 0.53${\pm}$0.28 0.79${\pm}$0.13 0.38${\pm}$0.22 -TestLayout
0.11${\pm}$0.03 0.13${\pm}$0.04 0.11${\pm}$0.03 0.15${\pm}$0.06
0.13${\pm}$0.06 0.17${\pm}$0.04 0.12${\pm}$0.06 -TestCountPlus 0.06${\pm}$0.03
0.07${\pm}$0.03 0.09${\pm}$0.05 0.08${\pm}$0.04 0.05${\pm}$0.02
0.08${\pm}$0.04 0.05${\pm}$0.02 -TestAll 0.05${\pm}$0.01 0.06${\pm}$0.03
0.08${\pm}$0.02 0.08${\pm}$0.02 0.06${\pm}$0.02 0.09${\pm}$0.04
0.06${\pm}$0.03 Average 0.42${\pm}$0.10 0.34${\pm}$0.12 0.47${\pm}$0.04
0.39${\pm}$0.12 0.30${\pm}$0.15 0.46${\pm}$0.06 0.22${\pm}$0.12 MoveToCorner-
Demo 0.67${\pm}$0.09 0.60${\pm}$0.09 0.67${\pm}$0.14 0.72${\pm}$0.15
0.78${\pm}$0.07* 0.63${\pm}$0.13 0.62${\pm}$0.14 -TestDynamics 0.59${\pm}$0.10
0.58${\pm}$0.10 0.66${\pm}$0.09 0.69${\pm}$0.11 0.76${\pm}$0.09*
0.65${\pm}$0.12 0.64${\pm}$0.16 -TestColour 0.35${\pm}$0.08 0.32${\pm}$0.16
0.39${\pm}$0.09 0.45${\pm}$0.21 0.51${\pm}$0.19 0.37${\pm}$0.10
0.57${\pm}$0.16* -TestShape 0.43${\pm}$0.10 0.43${\pm}$0.16 0.51${\pm}$0.09
0.52${\pm}$0.16 0.54${\pm}$0.14 0.49${\pm}$0.12 0.47${\pm}$0.15 -TestJitter
0.61${\pm}$0.10 0.54${\pm}$0.09 0.62${\pm}$0.12 0.63${\pm}$0.16
0.74${\pm}$0.08* 0.62${\pm}$0.11 0.61${\pm}$0.16 -TestAll 0.19${\pm}$0.09
0.21${\pm}$0.20 0.25${\pm}$0.14 0.31${\pm}$0.19 0.38${\pm}$0.18*
0.20${\pm}$0.13 0.39${\pm}$0.16* Average 0.48${\pm}$0.09 0.45${\pm}$0.10
0.52${\pm}$0.07 0.55${\pm}$0.15 0.62${\pm}$0.11* 0.49${\pm}$0.08
0.55${\pm}$0.14 MoveToRegion-Demo 0.96${\pm}$0.04 0.95${\pm}$0.04
0.95${\pm}$0.07 0.97${\pm}$0.03 0.96${\pm}$0.04 0.96${\pm}$0.08
0.89${\pm}$0.09 -TestDynamics 0.94${\pm}$0.05 0.94${\pm}$0.05 0.94${\pm}$0.07
0.97${\pm}$0.03 0.95${\pm}$0.04 0.95${\pm}$0.06 0.82${\pm}$0.20 -TestColour
0.65${\pm}$0.16 0.68${\pm}$0.10 0.69${\pm}$0.23 0.68${\pm}$0.12
0.68${\pm}$0.21 0.67${\pm}$0.19 0.57${\pm}$0.21 -TestJitter 0.92${\pm}$0.06
0.93${\pm}$0.05 0.96${\pm}$0.08 0.97${\pm}$0.04 0.97${\pm}$0.03
0.94${\pm}$0.08 0.69${\pm}$0.25 -TestLayout 0.59${\pm}$0.09 0.61${\pm}$0.10
0.59${\pm}$0.09 0.66${\pm}$0.06 0.64${\pm}$0.14 0.65${\pm}$0.16
0.39${\pm}$0.18 -TestAll 0.28${\pm}$0.09 0.36${\pm}$0.09 0.28${\pm}$0.08
0.29${\pm}$0.06 0.32${\pm}$0.08 0.35${\pm}$0.10 0.24${\pm}$0.09 Average
0.72${\pm}$0.07 0.74${\pm}$0.04 0.74${\pm}$0.06 0.76${\pm}$0.03
0.75${\pm}$0.07 0.75${\pm}$0.09 0.60${\pm}$0.14

Table A8:  Complete results for all variants of the evaluated MAGICAL tasks,
with GAIL plus RepL pretraining.

##  Appendix H Loss curves

###  H.1 Pretrain

###  H.2 Joint Training

##  Appendix I Additional contrastive learning ablations

Table A9 presents ablations for SimCLR on our three MAGICAL tasks, using BC as
the downstream IL algorithm. In particular, we experiment with:

Projection heads

By default, SimCLR uses a “symmetric” projection strategy that applies the
same projection head to the encoded contexts and encoded target before
computing the loss. We also try using asymmetric projection heads, which are
allowed to apply different transforms to the target and context embeddings,
and additionally experiment with removing projection heads entirely so that we
are computing SimCLR loss directly on the encoder representation.

Compression

We experiment with compression by replacing the default SimCLR loss with the
CEB loss [40], but leaving all other training and architecture details the
same.

Momentum

In the momentum ablation, we replace the SimCLR loss and encoder with a MoCo-
style [13] loss and momentum encoder. Again, all other training and
architecture details are left the same as in our SimCLR implementation.

In Table A9, we see that none of these modifications significantly improve
performance over standard SimCLR. For this reason, we expect that using
different projection heads, using compression, and using momentum are unlikely
to affect the conclusions of our work.

Task Asymm. proj. No proj. CEB loss Momentum SimCLR MatchRegions-Demo
0.43${\pm}$0.03 0.44${\pm}$0.02 0.42${\pm}$0.03 0.45${\pm}$0.04
0.42${\pm}$0.04 MoveToCorner-Demo 0.78${\pm}$0.03 0.83${\pm}$0.03
0.80${\pm}$0.03 0.83${\pm}$0.03 0.86${\pm}$0.06 MoveToRegion-Demo
0.82${\pm}$0.01 0.80${\pm}$0.03 0.83${\pm}$0.02 0.81${\pm}$0.01
0.81${\pm}$0.02

Table A9:  Ablations for SimCLR variants on MAGICAL. We used SimCLR as a
pretraining step for BC. Significance levels were evaluated relative to
vanilla SimCLR (the rightmost column) using a one-sided Welch’s t-test at
$p<0.05$, as with our other results. None of these results differ
significnatly from SimCLR, and so none are starred.