Teachable Reinforcement Learning via Advice Distillation.

Teachable Reinforcement Learning via
Advice Distillation
Olivia Watkins
UC Berkeley
oliviawatkins@berkeley.eduTrevor Darrell
UC Berkeley
trevor@eecs.berkeley.edu
Pieter Abbeel
UC Berkeley
pabbeel@cs.berkeley.eduJacob Andreas
MIT
jda@mit.eduAbhishek Gupta
UC Berkeley
abhigupta@berkeley.edu
Abstract
Training automated agents to complete complex tasks in interactive environments
is challenging: reinforcement learning requires careful hand-engineering of reward
functions, imitation learning requires specialized infrastructure and access to a
human expert, and learning from intermediate forms of supervision (like binary
preferences) is time-consuming and extracts little information from each human
intervention. Can we overcome these challenges by building agents that learn
from rich, interactive feedback instead? We propose a new supervision paradigm
for interactive learning based on “teachable” decision-making systems that learn
from structured advice provided by an external teacher. We begin by formalizing
a class of human-in-the-loop decision making problems in which multiple forms
of teacher-provided advice are available to a learner. We then describe a simple
learning algorithm for these problems that ﬁrst learns to interpret advice , then
learns from advice to complete tasks even in the absence of human supervision. In
puzzle-solving, navigation, and locomotion domains, we show that agents that learn
from advice can acquire new skills with signiﬁcantly less human supervision than
standard reinforcement learning algorithms and often less than imitation learning.
1 Introduction
Reinforcement learning (RL) offers a promising paradigm for building agents that can learn complex
behaviors from autonomous interaction and minimal human effort. In practice, however, signiﬁcant
human effort is required to design and compute the reward functions that enable successful RL [ 49]:
the reward functions underlying some of RL’s most prominent success stories involve signiﬁcant
domain expertise and elaborate instrumentation of the agent and environment [ 37,38,44,28,15].
Even with this complexity, a reward is ultimately no more than a scalar indicator of how good a
particular state is relative to others. Rewards provide limited information about how to perform tasks,
and reward-driven RL agents must perform signiﬁcant exploration and experimentation within an
environment to learn effectively. A number of alternative paradigms for interactively learning policies
have emerged as alternatives, such as imitation learning [ 40,20,50], DAgger [ 43], and preference
learning [ 10,6]. But these existing methods are either impractically low bandwidth (extracting little
information from each human intervention) [ 25,30,10] or require costly data collection [ 44,23]. It
has proven challenging to develop training methods that are simultaneously expressive and efﬁcient
enough to rapidly train agents to acquire novel skills.
Human learners, by contrast, leverage numerous, rich forms of supervision: joint attention [ 34],
physical corrections [ 5] and natural language instruction [ 9]. For human teachers, this kind of
35th Conference on Neural Information Processing Systems (NeurIPS 2021), virtual.
c(a) Grounding(b) Improvement(c) Evaluation
q(as,,c1)q(as,,c2)(as,)cccCoaching c1State sAction aReward rcturn rightc
cc0.30.10.2+1cgo straightc
cc0.20.20.10.2ccccTgt. coaching ci+1State sAction actravel to  waypoint  (3, 4)c
cccc
ccSrc. coaching cisimple coaching-conditional policycomplex coaching- conditional policyunconditional policy(as,)cTask ccmaze1maze2cTask ccmaze1maze2c
c[0.30.10.2]c
c[0.20.20.1](agent takes actions without advice)reinforcement learningadvice distillationadvice distillation
turn rightgo straight
0.30.10.20.20.20.1Figure 1: Three phases of teachable reinforcement learning. During the grounding phase (a), we train an
advice-conditional policy through RL q(a|s,t,c1)that can interpret a simple form of advice c1. During the
improvement phase (b), an external coach provides real-time coaching, which the agent uses to learn more
complex advice forms and ultimately an advice-independent policy p(a|s,t). During the evaluation phase, the
advice-independent policy p(a|s,t)is executed to accomplish a task without additional human feedback.
coaching is often no more costly to provide than scalar measures of success, but signiﬁcantly more
informative for learners. In this way, human learners use high-bandwidth, low-effort communication
as a means to ﬂexibly acquire new concepts or skills [ 46,33]. Importantly, the interpretation of some
of these feedback signals (like language), is itself learned, but can be bootstrapped from other forms
of communication: for example, the function of gesture and attention can be learned from intrinsic
rewards [39]; these in turn play a key role in language learning [31].
This paper proposes a framework for training automated agents using similarly rich interactive
supervision. For instance, given an agent learning a policy to navigate and manipulate objects in a
simulated multi-room object manipulation problem (e.g., Fig 3 left), we train agents using not just
reward signals but advice about what actions to take (“move left”), what waypoints to move towards
(“move towards (1,2)”), and what sub-goals to accomplish (“pick up the yellow ball”), offering
human supervisors a toolkit of rich feedback forms that direct and modify agent behavior. To do so,
we introduce a new formulation of interactive learning, the Coaching-Augmented Markov Decision
Process (CAMDP), which formalizes the problem of learning from a privileged supervisory signal
provided via an observation channel. We then describe an algorithmic framework for learning in
CAMDPs via alternating advice grounding anddistillation phases. During the grounding phase,
agents learn associations between teacher-provided advice and high-value actions in the environment;
during distillation, agents collect trajectories with grounded models and interactive advice, then
transfer information from these trajectories to fully autonomous policies that operate without advice.
This formulation allows supervisors to guide agent behavior interactively, while enabling agents to
internalize this guidance to continue performing tasks autonomously once the supervisor is no longer
present. Moreover, this procedure can be extended to enable bootstrapping of grounded models that
use increasingly sparse and abstract advice types, leveraging some types of feedback to ground others.
Experiments show that models trained via coaching can learn new tasks more efﬁciently and with
20x less human supervision than naïve methods for RL across puzzle-solving [ 8], navigation [ 14],
and locomotion domains [8].
In summary, this paper describes: (1) a general framework (CAMDPs) for human-in-the-loop RL with
rich interactive advice; (2) an algorithm for learning in CAMDPs with a single form of advice; (3) an
extension of this algorithm that enables bootstrapped learning of multiple advice types; and ﬁnally
(4) a set of empirical evaluations on discrete and continuous control problems in the BabyAI [ 8] and
D4RL [ 14] environments. It thus offers a groundwork for moving beyond reward signals in interactive
learning, and instead training agents with the full range of human communicative modalities.
2 Coaching Augmented Markov Decision Processes
To develop our procedure for learning from rich feedback, we begin by formalizing the environments
and tasks for which feedback is provided. This formalization builds on the framework of multi-task
RL and Markov decision processes (MDP), augmenting them with advice provided by a coach in the
loop through an arbitrary prescriptive channel of communication. Conider the grid-world environment
depicted in Fig 3 left [ 8].Tasks in this environment specify particular speciﬁc desired goal states;
e.g. “place the yellow ball in the green box and the blue key in the green box” or “open all doors in
2
the blue room.” In multi-task RL, a learner’s objective is produce a policy p(at|st,t)that maximizes
reward in expectation over tasks t. More formally, a multi-task MDP is deﬁned by a 7-tuple
M⌘(S,A,T,R,r(s0),g,p(t)), where Sdenotes the state space, Adenotes the action space,
p:S⇥A⇥S7!R0denotes the transition dynamics, r:S⇥A⇥t7!Rdenotes the reward
function, r:S7!R0denotes the initial state distribution, g2[0,1]denotes the discount factor and
p(t)denotes the distribution over tasks. The objective in a multi-task MDP is to learn a policy pqthat
maximizes the expected sum of discounted returns in expectation over tasks: max qJE(pq,p(t)) =
Eat⇠pq(·|st,t)
t⇠p(t)[Â•
t=0gtr(st,at,t)].
Why might additional supervision beyond the reward signal be useful for solving this optimization
problem? Suppose the agent in Fig 3 is in the (low-value) state shown in the ﬁgure, but could reach a
high-value state by going “right and up” towards the blue key. This fact is difﬁcult to communicate
through a scalar reward, which cannot convey information about alternative actions. A side channel
for providing this type of rich information at training-time would be greatly beneﬁcial.
We model this as follows: a coaching-augmented MDP (CAMDP) consists of an ordinary multi-
task MDP augmented with a set of coaching functions C={C1,C2,···,Ci}, where each Cj
provides a different form of feedback to the agent. Like a reward function, each coaching function
models a form of supervision provided externally to the agent (by a coach ); these functions may
produce informative outputs densely (at each timestep) or only infrequently. Unlike rewards, which
give agents feedback on the desirability of states and actions they have already experienced, this
coaching provides information about what the agent should do next.1As shown in Figure 3, advice
can take many forms, for instance action advice ( c0), waypoints ( c1), language sub-goals ( c2), or any
other local information relevant to task completion.2Coaching in a CAMDP is useful if it provides
an agent local guidance on how to proceed toward a goal that is inferrable from the agent’s current
observation, when the mapping from observations and goals to actions has not yet been learned.
As in standard reinforcement learning in an multi-task MDP, the goal in a CAMDP is to learn a policy
pq(·|st,t)that chooses an action based on Markovian state stand high level task information t
without interacting with cj. However, we allow learning algorithms to use the coaching signal cjto
learn this policy more efﬁciently at training time (although this is unavailable during deployment).
For instance, the agent in Fig 3 can leverage hints “go left” or “move towards the blue key” to guide
its exploration process but it eventually must learn how to perform the task without any coaching
required. Section 3 decribes an algorithm for acquiring this independent, multi-task policy pq(·|st,t)
from coaching feedback, and Section 4 presents an empirical evaluation of this algorithm.
3 Leveraging Advice via Distillation
3.1 Preliminaries
The challenge of learning in a CAMDP is twofold: ﬁrst, agents must learn to ground coaching signals
in concrete behavior; second, agents must learn from these coaching signals to independently solve
the task of interest in the absence of any human supervision. To accomplish this, we divide agent
training into three phases: (1) a grounding phase, (2) an improvement phase and (3) an evaluation
phase.
In the grounding phase, agents learn how to interpret coaching. The result of the grounding phase
is a surrogate policy q(at|st,t,c)that can effectively condition on coaching when it is provided in
the training loop. As we will discuss in Section 3.2, this phase can also make use of a bootstrapping
process in which more complex forms of feedback are learned using signals from simpler ones.
During the improvement phase, agents use the ability to interpret advice to learn new skills. Specif-
ically, the learner is presented with a novel task ttestthat was not provided during the grounding
phase, and must learn to perform this task using only a small amount of interaction in which advice
cis provided by a human supervisor who is present in the loop. This advice, combined with the
1While the design of optimal coaching strategies and explicit modeling of coaches are important research
topics [ 16], this paper assumes that the coach is ﬁxed and not explicitly modeled. Our empirical evaluation use
both scripted coaches and human-in-the-loop feedback.
2When only a single form of advice is available to the agent, we omit the superscript for clarity.
3
learned surrogate policy q(at|st,t,c), can be used to efﬁciently acquire an advice- independent policy
p(at|st,t), which can perform tasks without requiring any coaching.
Finally, in the evaluation phase, agent performance is evaluated on the task ttestby executing the
advice-independent, multi-task policy p(at|st,ttest)in the environment.
3.2 Grounding Phase: Learning to Interpret Advice
The goal of the grounding phase is to learn a mapping from advice to contextually appropriate actions,
so that advice can be used for quickly learning new tasks. In this phase, we run RL on a distribution
of training tasks p(t). As the purpose of these training environments is purely to ground coaching,
sometimes called “advice”, the tasks may be much simpler than test-time tasks. During this phase, the
agent uses access to a reward function r(s,a,c), as well as the advice c(s,a)to learn a surrogate policy
qf(a|s,t,c). The reward function r(s,a,c)is provided by the coach during the grounding phase only
and rewards the agent for correctly following the provided coaching, not just for accomplishing the
task. Since coaching instructions (e.g. cardinal directions) are much easier to follow than completing
a full task, grounding can be learned quickly. The process of grounding is no different than standard
multi-task RL, incorporating advice c(s,a)as another component of the observation space. This
formulation makes minimal assumptions about the form of the coaching c.
During this grounding process, the agent’s optimization objective is:
max
fJ(q)=E t⇠p(t)
at⇠qf(at|st,t,c)
Â
tr(st,at,c)
, (1)
Bootstrapping Multi-Level Advice The previous section described how to train an agent to in-
terpret a single form of advice c. In practice, a coach might ﬁnd it useful to use multiple forms of
advice—for instance high-level language sub-goals for easy stages of the task and low-level action
advice for more challenging parts of the task. While high-level advice can be very informative for
guiding the learning of new tasks in the improvement phase, it can often be quite difﬁcult to ground
quickly pure RL. Instead of relying on RL, we can bootstrap the process of grounding one form
of advice chfrom a policy q(a|s,t,cl)that can interpret a different form of advice cl. In particular,
we can use a surrogate policy which already understands (using the grounding scheme described
above) low-level advice q(a|s,t,cl)to bootstrap training of a surrogate policy which understands
higher-level advice q(a|s,t,ch). We call this process “bootstrap distillation”.
Intuitively, we use a supervisor in the loop to guide an advice-conditional policy that can interpret
a low-level form of advice qf1(a|s,t,cl)to perform a training task, obtaining trajectories D=
{(s0,a0,cl
0,ch
0),(s1,a1,cl
1,ch
1)···,(sH,aH,cl
H,ch
H)}N
j=1, then distilling the demonstrated behavior via
supervised learning into a policy qf2(a|s,t,ch)that can interpret higher-level advice to perform this(a) In-the-loop advice
ccc(s,a*,,c)
coached rollouts from conditional policydistillation into unconditional policy(b) Off-policy advice
ccc(s,a,)
uncoached rollouts from unconditional policydistillation into unconditional policyhindsight coaching and action relabeling
ccc(s,a*,,c)Figure 2: Illustration of the procedure of advice distillation in the on-policy and off-policy settings. During
on-policy advice distillation, the advice-conditional surrogate policy q(a|s,t,c)is coached to get optimal
trajectories. These trajectories are then distilled into an unconditional model p(a|s,t)using supervised learning.
During off-policy distillation, trajectories are collected by the unconditional policy and trajectories are relabeled
with advice after the fact. After this, we use the advice-conditional policy q(a|s,t,c)to relabel trajectories with
optimal actions. These trajectories can then be distilled into an unconditional policy.
4
new task without requiring the low level advice any longer. More speciﬁcally, we make use of an
input remapping solution, as seen in Levine et al. [28], where the policy conditioned on advice clis
used to generate optimal action labels, which are then remapped to observations with a different form
of advice chas input. To bootstrap the understanding of an abstract form of advice chfrom a more
low level one cl, the agent optimizes the following objective to bootstrap the agent’s understanding of
one advice type from another:
D={(s0,a0,cl
0,ch
0),(s1,a1,cl
1,ch
1),···,(sH,aH,cl
H,ch
H)}N
j=1
s0⇠p(s0),at⇠qf1(at|st,t,cl),st+1⇠p(st+1|st,at)
max
f2E(st,at,cht,t)⇠Dh
logqf2(at|st,t,ch
t)i
With this procedure, we only need to use RL to ground the simplest, fastest-learned advice form, and
we can use more efﬁcient bootstrapping to ground the others.
3.3 Improvement Phase: Learning New Tasks Efﬁciently with Advice
At the end of the grounding phase, we have an advice-following agent qf(a|s,t,c)that can interpret
various forms of advice. Ultimately, we want a policy p(a|s,t)which is able to succeed at performing
the new test task ttest,without requiring advice at evaluation time. To achieve this, we make use
of a similar idea to the one described above for bootstrap distillation. In the improvement phase,
we leverage a supervisor in the loop to guide an advice-conditional surrogate policy qf(a|s,t,c)to
perform the new task ttest, obtaining trajectories D={s0,a0,c0,s1,a1,c1,···,sH,aH,cH}N
j=1, then
distill this behavior into an advice-independent policy pq(a|s,t)via behavioral cloning. The result is
a policy trained using coaching, but ultimately able to select tasks even when no coaching is provided.
In Fig 3 left, this improvement process would involve a coach in the loop providing action advice or
language sub-goals to the agent during learning to coach it towards successfully accomplishing a
task, and then distilling this knowledge into a policy that can operate without seeing action advice or
sub-goals at execution time. More formally, the agent optimizes the following objective:
D={s0,a0,c0,s1,a1,c1,···,sH,aH,cH}N
j=1
s0⇠p(s0),at⇠qf(at|st,t,ct),st+1⇠p(st+1|st,at)
max
qE(st,at,t)⇠D[logpq(at|st,t)]
This improvement process, which we call advice distillation, is depicted Fig 2. This distillation
process is preferable over directly providing demonstrations because the advice provided can be
more convenient than providing an entire demonstration (for instance, compare the difﬁculty of
producing a demo by navigating an agent through an entire maze to providing a few sparse waypoints).
Interestingly, even if the new tasks being solved ttestare quite different from the training distribution
of tasks p(t), since advice c(for instance waypoints) is provided locally and is largely invariant to
this distribution shift, the agent’s understanding of advice generalizes well.
Learning with Off-Policy Advice One limitation to the improvement phase procedure described
above is that advice must be provided in real time. However, a small modiﬁcation to the algorithm
allows us to train with off-policy advice. During the improvement phase, we roll out an initially-
untrained advice-independent policy p(a|s,t). After the fact, the coach provides high-level advice
chat a multiple points along the trajectory. Next, we use the advice-conditional surrogate policy
qf(a|s,t,c)to relabel this trajectory with near-optimal actions at each timestep. This lets us use
behavioral cloning to update the advice-free agent on this trajectory. While this relabeling process
must be performed multiple times during training, it allows a human to coach an agent without
providing real-time advice , which can be more convenient. This process can be thought of as the
coach performing DAgger [ 42] at the level of high-level advice (as was done in in [ 26]) rather than
low-level actions. This procedure can be used for both the grounding and improvement phases.
Mathematically, the agent optimizes the following objective:
D={s0,a0,c0,s1,a1,c1,···,sH,aH,cH}N
j=1
s0⇠p(s0),at⇠p(at|st,t),st+1⇠p(st+1|st,at)
max
qE(st,t)⇠D
a⇤⇠qf(at|st,t,c)[logpq(a⇤|st,t)]
5
3.4 Evaluation Phase: Executing tasks Without a Supervisor
In the evaluation phase, the agent simply needs to be able to perform the test tasks ttestwithout
requiring a coach in the loop. We run the advice-independent agent learned in the improvement phase,
p(a|s,t)on the test task ttestand record the average success rate.
4 Experimental Evaluation
We aim to answer the following questions through our experimental evaluation (1) Can advice be
grounded through interaction with the environment via supervisor in the loop RL? (2) Can grounded
advice allow agents to learn new tasks more efﬁciently than standard RL? (3) Can agents bootstrap
the grounding of one form of advice from another?
4.1 Evaluation Domains
Instruction: Navigate to (x, y)“Pick up a blue key”Action advice: Waypoint:Subgoal:Action: TurnLeftWaypoint: (3, 7)“Go to the yellow door”Direction advice: Cardinal AdviceWaypoint:Direction: [.17, -.23]Direction: WestWaypoint (3, 4)
Navigate to (x, y)Direction advice: Cardinal Advice:Waypoint:Direction: [.17, -.23]Direction: WestWaypoint (3, 4)
Figure 3: Evaluation Domains. (Left) BabyAI (Middle) Point Maze Navigation (Right) Ant Navigation. The
associated task instructions are shown, as well as the types of advice available in each domain.
BabyAI: In the open-source BabyAI [ 8] grid-world, an agent is given tasks involving navigation,
pick and place, door-opening and multi-step manipulation. We provide three types of advice:
1.Action Advice: Direct supervision of the next action to take.
2.OffsetWaypoint Advice: A tuple (x, y, b), where (x, y) is the goal coordinate minus the
agent’s current position, and b tells the agent whether to interact with an object.
3.Subgoal Advice: A language subgoal such as “Open the blue door.”
2-D Maze Navigation (PM): In the 2D navigation environment, the goal is to reach a random target
within a procedurally generated maze. We provide the agent different types of advice:
1.Direction Advice : The vector direction the agent should head in.
2.Cardinal Advice : Which of the cardinal directions (N, S, E, W) the agent should head in.
3.Waypoint Advice : The (x,y) position of a coordinate along the agent’s route.
4.OffsetWaypoint Advice : The (x,y) waypoint minus the agent’s current position.
Ant-Maze Navigation (Ant): The open-source ant-maze navigation domain [ 14] replaces the simple
point mass agent with a quadrupedal “ant” robot. The forms of advice are the same as the ones
described above for the point navigation domain.
In all domains, we describe advice forms provided each timestep (Action Advice and Direction
Advice) as “low-level” advice, and advice provided less frequently as “high-level” advice. We present
experiments involving both scripted coaches and real human-in-the-loop advice.
6
4.2 Experimental Setup
For the environments listed above, we evaluate the ability of the agent to perform grounding efﬁciently
on a set of training tasks, to learn new test tasks quickly via advice distillation, and to leverage one
form of advice to bootstrap another. The details of the exact set of training and testing tasks, as well
as architecture and algorithmic details, are provided in the appendix.
We evaluate all the environments using the metric of advice efﬁciency rather than sample efﬁciency.
By advice efﬁciency, we are evaluating the number of instances of coach-in-the-loop advice that are
needed in order to learn a task. In real-world learning tasks, this coach is typically a human, and the
cost of training largely comes from the provision of supervision (rather than time the agent spends
interacting with the environment). The same is true for other forms of supervision such as behavioral
cloning and RL (unless the human spends extensive time instrumenting the environment to allow
autonomous rewards and resets). This “advice units” metric more accurately reﬂects the true quantity
we would like to measure: the amount of human time and effort needed to provide a particular course
of coaching. For simplicity, we consider every time a supervisor provides any supervision, such as
a piece of advice or a scalar reward, to constitute one advice unit . We measure efﬁciency in terms
of how many advice units are needed to learn a task. We emphasize that this metric makes a strong
simplifying assumption—that all forms of advice have the same cost—which is certainly not true
for real-world supervision. However, it is challenging to design a metric which accurately captures
human effort. In Section 4.7 we validate our method by measuring the real human interaction time
needed to train agents. We also plot more traditional sample efﬁciency measures in Appendix D.
We compare our proposed framework to an RL baseline that is provided with a task instruction but
no advice. In the improvement phase, we also compare with behavioral cloning from an expert for
environments where it is feasible to construct an oracle.
4.3 Grounding Prescriptive Advice during Training
Figure 4: Left: Performance during the grounding phase (Section 3.2). All curves are trained with shaped-rewardRL. We compare agents which condition on high-level advice (shades of blue) to ones with access to low-leveladvice (red) to an advice-free baseline (gray). Takeaways: (a) the agent is able to ground advice, which suggeststhat our advice-conditional policy may be useful for coaching; (b) grounding certain high-level advice formsthrough RL is slow, which is why bootstrapping is necessary. Right: Bootstrapping is able to quickly use existinggrounded advice forms (OffsetWaypoint for Point Maze and Ant Maze envs, ActionAdvice for BabyAI) toground additional forms of advice.Fig 4 shows the results of the grounding phase, where the agent grounds advice by training anadvice-conditional policy through RL. We observe the the agent learns the task more quickly whenprovided with advice, indicating that the agent is learning to interpret advice to complete tasks.However, we also see that the agent fails to improve much when conditioning on some more abstractforms of advice, such as waypoint advice in the ant environment. This indicates that the advice formhas not been grounded properly through RL. In cases like this, we instead must instead ground theseadvice forms through bootstrapping, as discussed in Section 3.2.4.4 Bootstrapping Multi-Level FeedbackOnce we have successfully grounded the easiest form of advice, in each environment, we efﬁcientlyground the other forms using the bootstrapping procedure from Section 3.2. As we see in Fig 4,bootstrap distillation is able to ground new forms of advice signiﬁcantly more efﬁciently than if westart grounding things from scratch with naïve RL. It performs exceptionally well even for adviceforms where naïve RL does not succeed at all, while providing additional speed up for environments7
where it does. This suggests that advice is not just a tool to solve new tasks, but also a tool for
grounding more complex forms of communication for the agent.
4.5 Learning New Tasks with Grounded Prescriptive Advice
Point MazeDirectionCardinalWaypointOffsetRLOracle6x6 Maze 0.9±0.020.95±0.05 0.99±0.01 0.99±0.010.27±0.01 0.87±0.017x10 Maze0.75±0.090.77±0.060.74±0.09 0.9±0.050.09±0.04 0.73±0.0510x10 Maze 0.69±0.06 0.67±0.04 0.62±0.040.85±0.040.11±0.04 0.64±0.0613x13 Maze 0.16±0.040.35±0.080.22±0.050.45±0.030.08±0.04 0.28±0.04Ant MazeDirectionCardinalWaypointOffsetRL3x3 Maze 0.25±0.17 0.38±0.20.77±0.20.8±0.210.0±0.06x6 Maze 0.04±0.04 0.32±0.110.56±0.25 0.55±0.250.0±0.07x10 Maze0.0±0.00.0±0.00.0±0.00.0±0.00.0±0.0BabyAIAction AdviceOffsetWaypointSubgoalRLOracle BCTest Env 10.31±0.15 0.51±0.14 0.53±0.150.0±0.00.31±0.14Test Env 20.53±0.16 0.66±0.16 0.43±0.170.18±0.070.6±0.06Test Env 3 0.14±0.010.2±0.06 0.2±0.080.04±0.030.16±0.04Test Env 4 0.04±0.010.1±0.02 0.1±0.050.0±0.00.04±0.03Test Env 5 0.07±0.03 0.13±0.020.2±0.110.0±0.00.05±0.02Test Env 60.44±0.10.48±0.090.28±0.02 0.17±0.090.43±0.12Test Env 7 0.32±0.04 0.42±0.060.54±0.120.01±0.01 0.26±0.03Figure 5: Learning new tasks through distillation. The agent uses an already-grounded advice channel to performthe distillation process from Section 3.3 to train an advice-free agent. Results show the success rate of theadvice-free new agent. Left, we show representative curves for a few environments. Colors designate supervisionused: shades of blue = high level advice; red = low level advice; black = oracle demonstrations; gray = shapedrewards. Right: We show success rates (mean, std) over 3 seeds for a larger set of environments. Runs are boldedif std intervals overlapped with the highest mean. Success rates are evaluated at3e5steps for Point Maze andAnt Maze and5e5steps for BabyAI. Takeaway: once advice is grounded, in general it is most efﬁcient to teachthe agents new tasks by providing high-advice. There are occasional exceptions, discussed in Appendix G.
Figure 6: “Best advice” is OffsetAdvice. Y-axis includesadvice from both grounding and improvement across allfour Point Maze test envs. RL results stretch off the plot,indicating we were unable to run RL for long enough toconverge to the success rates of the other methods.Finally, we evaluate whether we can usegrounded advice to guide the agent through newtasks. In most cases, we directly used advice-conditional policies learned during groundingand bootstrapping. However, about half of theBabyAI high-level advice policies performedpoorly on the test environments. In this case, weﬁnetuned the policies with a few (<4k) samplescollected with rollouts from a lower-level bettergrounded advice form.As we can see in Fig 5, agents which are trainedthrough distillation from an abstract coach onaverage train with less supervision than RL agents. Providing high-level advice can even sometimesoutperform providing demonstrations, as the high-level advice allows the human to coach the agentthrough a successful trajectory without needing to provide an action at each timestep. It is about asefﬁcient to provide low-level advice as to provide demos (when demos are available), as both involveproviding one supervision unit per timestep.Advice grounding on the new tasks is not always perfect, however. For Instance, in BabyAI TestEnv 2 in Figure 5, occasional errors in the advice-conditional policy’s interpretation of high-adviceresult in it being just as efﬁcient efﬁcient to provide low-level advice or demos as it is to providehigh-level advice (though both are more efﬁcient than RL). When grounding is poor, the convergedﬁnal policy may not be fully successful. Baseline methods, in contrast, may ultimately convergeto higher rates, even if they take far more samples. For instance, RL never succeeds in AntMaze3x3 and 6x6 in the plots in Figure 5, but if training is continued for 1e6 advice units, RL achievesnear-perfect performance, whereas our method plateaus. This suggests our method is most usefulwhen costly supervision is the main constraint.8
The curve in Figure 5 is not entirely a fair comparison - after all, we are not taking into account
the advice units used to train the advice-conditional surrogate policy. However, it’s also not fair to
include this cost for each test env, since the up-front cost of grounding advice gets amortized over
a large set of downstream tasks. Figure 6 summarizes the total number of samples needed to train
each model to convergence on the Point Maze test environments, including all supervision provided
during grounding and improvement. We see that when we use the best advice form, our method is
8x more efﬁcient than demos, and over 20x more efﬁcient than dense-reward RL. In the PointMaze
environment, the cost of grounding becomes worthwhile with only 4 test envs. In other environments
such as Ant, it may take many more test envs than the three we tested on. This suggests that our
method is most appropriate when the agent will be used on a large set of downstream tasks.
4.6 Off-Policy Advice Relabeling
One limitation of the improvement phase as described Section 4.5 is that the human coach has to
be continuously present as the agent is training to provide advice on every trajectory. We relax this
requirement by providing the advice in hindsight rather than in-the-loop using the procedure from
Section 3.3. Results are shown in Figure 7. IN the Point Maze and Ang envs, this DAgger-like
scheme for soliciting advice performs greater than or equal to real-time advice. However, it performs
worse in the BabyAI environment. In future work we will explore this approach further, as it removes
the need for a human to be constantly present in the loop and opens avenues for using active learning
techniques to label only the most informative trajectories.
4.7 Real Human Experiments
Figure 7: All curves show the success rate of an advice-free policy
trained via distillation from an advice-conditional surrogate policy.
All curves use the OffsetWaypoint advice form, and results are aver-
aged over three seeds. Takeaway: DAgger performs well on some
environments (Point Maze, Ant) but poorly on others (BabyAI).To validate the automated evaluation
above (and determine whether our
“advice unit” metric is a good proxy
for human effort), we performed an
additional set of experiments with
human-in-the-loop coaches. Advice-
conditional surrogate policies were
pre-trained to follow advice using
a scripted coach. The coaches (all
researchers at U.C. Berkeley) then
coached these agents through solving
new, more complex test environments.
Afterwards, an an advice-free policy
was distilled from the successful tra-
jectories. Humans provided advice
through a click interface. (For instance, they could click on the screen to provide a.) See Fig 8.
In the BabyAI environment, we provide OffsetWaypoint advice and compare against a behavioral
cloning (BC) baseline where the human provided per-timestep demonstrations using arrow keys. Our
method’s is higher variance and has a slightly lower mean success rate, but results are still largely
consistent with Figure 5, which showed that for the BabyAI env BC is competitive with our method.
In the Ant environment, demonstrations aren’t possible, and the agent does not explore well enough
to learn from sparse rewards. We compare against the performance of an agent coached by a scripted
coach providing dense, shaped rewards. We see that the agent trained with 30 minutes of coaching
by humans performs comparably to an RL agent trained with 3k more advice units.
5 Related Work
The learning problem studied in this paper belongs to a more general class of human-in-the-loop RL
problems [ 1,25,30,47,12]. Existing frameworks like TAMER [ 25,45] and COACH [ 30,4] also
use interactive feedback to train policies, but are restricted to scalar or binary rewards. In contrast,
our work formalizes the problem of learning from arbitrarily complex feedback signals. A distinct
line of work looks to learn how to perform tasks from binary feedback with human preferences,
for example by indicating which of two trajectory snippets a human might prefer [ 10,21,47,27].
9
These techniques receive only a single bit of information with every human interaction, making
human supervision time-consuming and tedious. In contrast, the learning algorithm we describe uses
higher-bandwidth feedback signals like language-based subgoals and directional nudges, provided
sparsely, to reduce the required effort from a supervisor.
Figure 8: Left, Middle: We compare the success of an advice-freepolicy trained in two test envs with real human coaching to a RLpolicy trained with a scripted reward. “RL 10x” means RL policyreceived 10x more advice units (left) or samples (middle). Right:success of advice-free policies trained with 30 mins of human time.Humans either coach the agent with our method or provide demos.Sample sizes are n=2 per condition for Ant, n=3 per condition forBabyAI, so the results are suggestive not conclusive.Learning from feedback, especiallyprovided in the form of natural lan-guage, is closely related to instructionfollowing in natural language process-ing [7,3,32,41]. In instruction fol-lowing problems, the goal is to pro-duce aninstruction-conditionalpolicythat can generalize to new natural lan-guage speciﬁcations of behavior (atthe level of either goals or action se-quences [24] and held-out environ-ments. Here, our goal is to produceanunconditionalpolicy that achievesgood task success autonomously—weuse instruction following models to in-terpret interactive feedback and scaf-fold the learning of these autonomouspolicies. Moreover, the advice pro-vided is not limited to task-level spec-iﬁcations, but instead allows for real-time, local guidance of behavior. Thisprovides signiﬁcantly greater ﬂexibil-ity in altering agent behavior.The use of language at training time to scaffold learning has been studied in several more speciﬁcsettings [29]: Co-Reyes et al.[11]describe a procedure for learning to execute ﬁxed target trajectoriesvia interactive corrections, Andreas et al.[2]use language to produce policyrepresentationsusefulfor reinforcement learning, while Jiang et al.[22]and Hu et al.[18]use language to guide the learningof hierarchical policies. Eisenstein et al.[13]and Narasimhan et al.[35]use side information fromlanguage to communicate information about environment dynamics rather than high-value actionsequences. In contrast to these settings, we aim to use interactive human in the loop advice to learnpolicies that can autonomously perform novel tasks, even when a human supervisor is not present.6 DiscussionSummary:In this work, we introduced a new paradigm for teacher in the loop RL, which we refer toas coaching augmented MDPs. We show that CAMPDs cover a wide range of scenarios and introducea novel framework to learn how to interpret and utilize advice in CAMDPs. We show that doing sohas the dual beneﬁts of being able to learn new tasks more efﬁciently in terms of human effortandbeing able to bootstrap one form of advice off of another for more efﬁcient grounding.Limitations:Our method relies on accurate grounding of advice, which does not always happen inthe presence of other correlated environment features (e.g. the advice to “open the door,” and thepresence of a door in front of the agent). Furthermore, while our method is more efﬁcient than BC orRL, it still requires signiﬁcant human effort. These limitations are discussed further in Appendix G.Societal impacts:As human in the loop systems such as the one described here are scaled upto real homes, privacy becomes a major concern. If we have learning systems operating aroundhumans, sharing data and incorporating human feedback into their learning processes, they need to becareful about not divulging private information. Moreover, human in the loop systems are constantlyoperating around humans and need to be especially safe.Acknowledgments:Thanks to experiment volunteers Yuqing Du, Kimin Lee, Anika Ramachandran,Philippe Hansen-Estruch, Alejandro Escontrela, Michael Chang, Sam Toyer, Ajay Jain, Dhruv Shah,Homer Walke. Funding by NSF GRFP and DARPA’s XAI, LwLL, and/or SemaFor program, as wellas BAIR’s industrial alliance programs.10
References
[1]D. Abel, J. Salvatier, A. Stuhlmüller, and O. Evans. Agent-agnostic human-in-the-loop reinforce-
ment learning. CoRR , abs/1701.04079, 2017. URL http://arxiv.org/abs/1701.04079 .
[2]J. Andreas, D. Klein, and S. Levine. Learning with latent language. In M. A. Walker, H. Ji, and
A. Stent, editors, NAACL , 2018.
[3]Y. Artzi and L. Zettlemoyer. Weakly supervised learning of semantic parsers for mapping
instructions to actions. Trans. Assoc. Comput. Linguistics , 1:49–62, 2013. URL https:
//tacl2013.cs.columbia.edu/ojs/index.php/tacl/article/view/27 .
[4]D. Arumugam, J. K. Lee, S. Saskin, and M. L. Littman. Deep reinforcement learning from
policy-dependent human feedback. CoRR , abs/1902.04257, 2019. URL http://arxiv.org/
abs/1902.04257 .
[5]A. Bajcsy, D. P. Losey, M. K. O’Malley, and A. D. Dragan. Learning robot objectives from
physical human interaction. In Conference on Robot Learning (CoRL) , 2017.
[6]D. S. Brown, W. Goo, and S. Niekum. Better-than-demonstrator imitation learning via
automatically-ranked demonstrations. In Conference on Robot Learning (CoRL) , 2019.
[7]D. L. Chen and R. J. Mooney. Learning to interpret natural language navigation instructions
from observations. In W. Burgard and D. Roth, editors, AAAI , 2011.
[8]M. Chevalier-Boisvert, D. Bahdanau, S. Lahlou, L. Willems, C. Saharia, T. H. Nguyen, and
Y. Bengio. Babyai: A platform to study the sample efﬁciency of grounded language learning.
InICLR , 2019.
[9]S. Chopra, M. H. Tessler, and N. D. Goodman. The ﬁrst crank of the cultural ratchet: Learning
and transmitting concepts through language. In A. K. Goel, C. M. Seifert, and C. Freksa, editors,
CogSci , 2019.
[10] P. F. Christiano, J. Leike, T. B. Brown, M. Martic, S. Legg, and D. Amodei. Deep reinforcement
learning from human preferences. In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach,
R. Fergus, S. V . N. Vishwanathan, and R. Garnett, editors, NeurIPS , 2017.
[11] J. D. Co-Reyes, A. Gupta, S. Sanjeev, N. Altieri, J. Andreas, J. DeNero, P. Abbeel, and S. Levine.
Guiding policies with language via meta-learning. In ICLR , 2019.
[12] C. A. Cruz and T. Igarashi. A survey on interactive reinforcement learning: Design principles
and open challenges. In R. Wakkary, K. Andersen, W. Odom, A. Desjardins, and M. G.
Petersen, editors, DIS ’20: Designing Interactive Systems Conference 2020, Eindhoven, The
Netherlands, July 6-10, 2020 , pages 1195–1209. ACM, 2020. doi: 10.1145/3357236.3395525.
URL https://doi.org/10.1145/3357236.3395525 .
[13] J. Eisenstein, J. Clarke, D. Goldwasser, and D. Roth. Reading to learn: Constructing features
from semantic abstracts. In EMNL , 2009.
[14] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4RL: datasets for deep data-driven
reinforcement learning. CoRR , abs/2004.07219, 2020. URL https://arxiv.org/abs/2004.
07219 .
[15] A. Gupta, J. Yu, T. Z. Zhao, V . Kumar, A. Rovinsky, K. Xu, T. Devlin, and S. Levine. Reset-free
reinforcement learning via multi-task learning: Learning dexterous manipulation behaviors
without human intervention. arXiv preprint arXiv:2104.11203 , 2021.
[16] D. Hadﬁeld-Menell, A. Dragan, P. Abbeel, and S. Russell. Cooperative inverse reinforcement
learning. arXiv preprint arXiv:1606.03137 , 2016.
[17] D. Hejna, L. Pinto, and P. Abbeel. Hierarchically decoupled imitation for morphological transfer.
InInternational Conference on Machine Learning , pages 4159–4171. PMLR, 2020.
11
[18] H. Hu, D. Yarats, Q. Gong, Y . Tian, and M. Lewis. Hierarchical decision making by generating
and following natural language instructions. In H. M. Wallach, H. Larochelle, A. Beygelzimer,
F. d’Alché-Buc, E. B. Fox, and R. Garnett, editors, NeurIPS , 2019.
[19] D. Y .-T. Hui, M. Chevalier-Boisvert, D. Bahdanau, and Y. Bengio. Babyai 1.1. arXiv preprint
arXiv:2007.12770 , 2020.
[20] A. Hussein, M. M. Gaber, E. Elyan, and C. Jayne. Imitation learning: A survey of learning
methods. ACM Comput. Surv. , 50(2):21:1–21:35, 2017. doi: 10.1145/3054912. URL https:
//doi.org/10.1145/3054912 .
[21] B. Ibarz, J. Leike, T. Pohlen, G. Irving, S. Legg, and D. Amodei. Reward learning from
human preferences and demonstrations in atari. In S. Bengio, H. M. Wallach, H. Larochelle,
K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, NeurIPS , 2018.
[22] Y. Jiang, S. Gu, K. Murphy, and C. Finn. Language as an abstraction for hierarchical deep
reinforcement learning. In NeurIPS , 2019.
[23] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, E. Jang, D. Quillen, E. Holly,
M. Kalakrishnan, V. Vanhoucke, and S. Levine. Qt-opt: Scalable deep reinforcement learn-
ing for vision-based robotic manipulation. CoRR , abs/1806.10293, 2018. URL http:
//arxiv.org/abs/1806.10293 .
[24] S. Karamcheti, E. C. Williams, D. Arumugam, M. Rhee, N. Gopalan, L. L. S. Wong, and
S. Tellex. A tale of two draggns: A hybrid approach for interpreting action-oriented and
goal-oriented instructions. In M. Bansal, C. Matuszek, J. Andreas, Y . Artzi, and Y . Bisk, editors,
RoboNLP@ACL , 2017.
[25] W. B. Knox and P. Stone. TAMER: Training an Agent Manually via Evaluative Reinforcement.
InIEEE 7th International Conference on Development and Learning , August 2008.
[26] H. Le, N. Jiang, A. Agarwal, M. Dudík, Y. Yue, and H. Daumé. Hierarchical imitation and
reinforcement learning. In International conference on machine learning , pages 2917–2926.
PMLR, 2018.
[27] K. Lee, L. Smith, and P. Abbeel. Pebble: Feedback-efﬁcient interactive reinforcement learning
via relabeling experience and unsupervised pre-training. In International Conference on
Machine Learning , 2021.
[28] S. Levine, C. Finn, T. Darrell, and P. Abbeel. End-to-end training of deep visuomotor policies.
The Journal of Machine Learning Research , 17(1):1334–1373, 2016.
[29] J. Luketina, N. Nardelli, G. Farquhar, J. N. Foerster, J. Andreas, E. Grefenstette, S. Whiteson,
and T. Rocktäschel. A survey of reinforcement learning informed by natural language. In IJCAI ,
2019.
[30] J. MacGlashan, M. K. Ho, R. T. Loftin, B. Peng, G. Wang, D. L. Roberts, M. E. Taylor, and
M. L. Littman. Interactive learning from policy-dependent human feedback. In ICML , 2017.
[31] N. M. McNeil, M. W. Alibali, and J. L. Evans. The role of gesture in children’s comprehension
of spoken language:now they need it, now they don’t. Journal of Nonverbal Behavior , 24
(2):131–150, 2000. doi: 10.1023/A:1006657929803. URL https://doi.org/10.1023/A:
1006657929803 .
[32] H. Mei, M. Bansal, and M. R. Walter. Listen, attend, and walk: Neural mapping of navigational
instructions to action sequences. In D. Schuurmans and M. P. Wellman, editors, AAAI , 2016.
[33] T. J. H. Morgan, N. T. Uomini, L. E. Rendell, L. Chouinard-Thuly, S. E. Street, H. M. Lewis, C. P.
Cross, C. Evans, R. Kearney, I. de la Torre, A. Whiten, and K. N. Laland. Experimental evidence
for the co-evolution of hominin tool-making teaching and language. Nature Communications ,6
(1):6029, 2015. doi: 10.1038/ncomms7029. URL https://doi.org/10.1038/ncomms7029 .
[34] P. Mundy and W. Jarrold. Infant joint attention, neural networks and social cognition. Neural
Networks , 23(8-9):985–997, 2010. doi: 10.1016/j.neunet.2010.08.009. URL https://doi.
org/10.1016/j.neunet.2010.08.009 .
12
[35] K. Narasimhan, R. Barzilay, and T. S. Jaakkola. Deep transfer in reinforcement learning by
language grounding. CoRR , abs/1708.00133, 2017. URL http://arxiv.org/abs/1708.
00133 .
[36] K. Nguyen, D. Misra, R. Schapire, M. Dudík, and P. Shafto. Interactive learning from activity
description. arXiv preprint arXiv:2102.07024 , 2021.
[37] OpenAI. Openai ﬁve. arxiv , 2018.
[38] OpenAI, I. Akkaya, M. Andrychowicz, M. Chociej, M. Litwin, B. McGrew, A. Petron, A. Paino,
M. Plappert, G. Powell, R. Ribas, J. Schneider, N. Tezak, J. Tworek, P. Welinder, L. Weng,
Q. Yuan, W. Zaremba, and L. Zhang. Solving rubik’s cube with a robot hand. CoRR ,
abs/1910.07113, 2019. URL http://arxiv.org/abs/1910.07113 .
[39] F. Poli, G. Serino, R. B. Mars, and S. Hunnius. Infants tailor their attention to maximize
learning. Science Advances , 6(39), 2020. doi: 10.1126/sciadv.abb5053. URL https://
advances.sciencemag.org/content/6/39/eabb5053 .
[40] D. Pomerleau. ALVINN: an autonomous land vehicle in a neural network. In D. S. Touretzky,
editor, NeurIPS , 1988.
[41] J. Roh, C. Paxton, A. Pronobis, A. Farhadi, and D. Fox. Conditional driving from natural
language instructions. In L. P. Kaelbling, D. Kragic, and K. Sugiura, editors, CoRL , 2019.
[42] S. Ross, G. Gordon, and D. Bagnell. A reduction of imitation learning and structured prediction
to no-regret online learning. In Proceedings of the fourteenth international conference on artiﬁ-
cial intelligence and statistics , pages 627–635. JMLR Workshop and Conference Proceedings,
2011.
[43] S. Ross, G. J. Gordon, and D. Bagnell. A reduction of imitation learning and structured
prediction to no-regret online learning. In G. J. Gordon, D. B. Dunson, and M. Dudík, editors,
AISTATS , 2011.
[44] J. Schulman, S. Levine, P. Abbeel, M. I. Jordan, and P. Moritz. Trust region policy optimization.
InICML , 2015.
[45] G. Warnell, N. R. Waytowich, V. Lawhern, and P. Stone. Deep TAMER: interactive agent
shaping in high-dimensional state spaces. In AAAI , 2018.
[46] S. R. Waxman and D. B. Markow. Words as invitations to form categories: evidence from 12-
to 13-month-old infants. Cogn Psychol , 29(3):257–302, Dec 1995.
[47] R. Zhang, F. Torabi, L. Guan, D. H. Ballard, and P. Stone. Leveraging human guidance for deep
reinforcement learning tasks. In S. Kraus, editor, IJCAI , 2019.
[48] F. Zhu, Y. Zhu, X. Chang, and X. Liang. Vision-language navigation with self-supervised
auxiliary reasoning tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision
and Pattern Recognition , pages 10012–10022, 2020.
[49] H. Zhu, J. Yu, A. Gupta, D. Shah, K. Hartikainen, A. Singh, V. Kumar, and S. Levine. The
ingredients of real world robotic reinforcement learning. In International Conference on
Learning Representations , 2020.
[50] B. D. Ziebart, A. L. Maas, J. A. Bagnell, and A. K. Dey. Maximum entropy inverse reinforcement
learning. In D. Fox and C. P. Gomes, editors, AAAI , 2008.
Checklist
1.For all authors...
(a)Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope? [Yes]
(b)Did you describe the limitations of your work? [Yes] See Section 6
13
(c)Did you discuss any potential negative societal impacts of your work? [Yes] See
Section 6
(d)Have you read the ethics review guidelines and ensured that your paper conforms to
them? [Yes] This work does not actually use human subjects, and is largely done in
simulation. But we have included a discussion in Section 6
2.If you are including theoretical results...
(a)Did you state the full set of assumptions of all theoretical results? [N/A] Math is used
as a theory/formalism, but we don’t make any provable claims about it.
(b)Did you include complete proofs of all theoretical results? [N/A]
3.If you ran experiments...
(a)Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)? [Yes] See Appendix
A for link to URL and run instructions in the README in the github repo.
(b)Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] See Appendix A.
(c)Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)? [Yes] All plots were created with 3 random seeds with std error
bars.
(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] See Appendix A
4.If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a)If your work uses existing assets, did you cite the creators? [Yes] Envs we used are
cited in section 4.1
(b)Did you mention the license of the assets? [Yes] This is in Appendix B
(c)Did you include any new assets either in the supplemental material or as a URL? [Yes]
We published the code and included all environments and assets as a part of this
(d)Did you discuss whether and how consent was obtained from people whose data you’re
using/curating? [Yes] We used three open source domains and collected our own data
on these domains.
(e)Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content? [N/A]
5.If you used crowdsourcing or conducted research with human subjects...
(a)Did you include the full text of instructions given to participants and screenshots, if
applicable? [No] We did not include full text since we didn’t use an exact script, but
we summarized the instructions and included images of the environments used.
(b)Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A] Only human involvement was data
collection with our system.
(c)Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation? [N/A] Human testers were volunteers
14

Olivia Watkins
Abhishek Gupta
Trevor Darrell
Pieter Abbeel
Jacob Andreas


An Empirical Investigation of Representation Learning for Imitation.

An Empirical Investigation of
Representation Learning for Imitation
Xin Chen
The University of Hong Kong
cyn0531@connect.hku.hkSam Toyer
UC Berkeley
sdt@berkeley.eduCody Wild
UC Berkeley
codywild@berkeley.edu
Scott Emmons
UC BerkeleyIan Fischer
Google ResearchKuang-Huei Lee
Google ResearchNeel Alex
UC BerkeleySteven Wang
UC Berkeley
Ping Luo
The University of Hong KongStuart Russell
UC BerkeleyPieter Abbeel
UC BerkeleyRohin Shah
UC Berkeley
Abstract
Imitation learning often needs a large demonstration set in order to handle the
full range of situations that an agent might ﬁnd itself in during deployment. How-
ever, 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-speciﬁc data.
Our Empirical Investigation of Representation Learning for Imitation (EIRLI)
investigates whether similar beneﬁts 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 ﬁnd that existing algo-
rithms for image-based representation learning provide limited value relative to a
well-tuned baseline with image augmentations. To explain this result, we investi-
gate differences between imitation learning and other settings where representation
learning hasprovided signiﬁcant beneﬁt, such as image classiﬁcation. Finally,
we release a well-documented codebase which both replicates our ﬁndings 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
Equal contribution, corresponding authors
35th Conference on Neural Information Processing Systems (NeurIPS 2021) Track on Datasets and Benchmarks.
Decoder
 Encoder
 Augmenter
Context, target
RepL
lossDataset
for repL
State, expert
action
Expert
demos Representation Learning (RepL)
Imitation Learning (IL)
Policy
IL
loss
Augmenter
 Encoder
Figure 1: A framework for the use of representation learning (RepL) in imitation learning. In the
pretraining setting, we ﬁrst train the encoder with RepL, then ﬁnetune end-to-end with IL. In the joint
training setting, the RepL objective is used as an auxiliary loss throughout IL training.
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 ﬁnd that, on average, RepL methods do signiﬁcantly outperform vanilla BC, but this beneﬁt 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 ﬁrst 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.We identify meaningful axes of variation in representation learning algorithm design, allowing us
to construct a modular framework to conceptually analyze these designs.
2.We use this framework to build a well documented, modular, and extensible code base, which we
release at github.com/HumanCompatibleAI/eirli .
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 augmen-
tations. 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
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 classiﬁcation 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]
2
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
V AE [9] Reconstruction 7 7 7 Pre
AugMix [10] Consistency 3 7 7 Joint
FixMatch [11] Consistency 3 7 7 Joint
CPC [12] Contrastive 7 7 3 Pre
MoCo [13] Contrastive 3 3 7 Pre
SimCLR [2] Contrastive 3 7 3 Pre
SimCLRv2 [14] Contrastive 3 7 3 Pre
BYOL [15] Bootstrap 3 3 3 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 otot;ot+1 3 7 7 7 7 P
PlaNet [17] R MPC + CEM o1:tot+1:T;rt+1:T3 7 7 7 7 J
CURL [18] C SAC otot 7 3 7 7 7 J
PI-SAC [3] C SAC otot+k;rt+k3 3 3 3 3 J
ATC [19] C SAC, PPO otot+k 7 3 3 3 7 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 xandy. Since the learning
signal derives from the relationship between xandy, the choice of xandythus has a signiﬁcant
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 classiﬁcation, 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
(V AE) [ 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 V AE, except that for a
context observation ot, we set the target to be a future observation ot+k. Now, the input observation
ot(context) is encoded into a vector representation and then decoded back into pixels, which is then
compared against the future observationot+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 rtto 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:
3
Reconstruction. Here, the goal is to reconstruct the target yfrom the representation z. Both the V AE
and temporal V AE in the previous section use a reconstructive loss, in which a decoded image d(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 (x1;y1);(x2;y2);:::; (xK;yK)
and use the same network to encode both the context and target into latent representations zie(zj
xi)andz0
ie(zjyi). A contrastive loss then incentivizes ziandz0
ito be similar to each other, but
different from zjandz0
jfor all other pairs j6=i. Typically, the contrastive loss function is chosen to
maximize the mutual information I(z;y), such as with the InfoNCE loss function [21]:
LInfoNCE =E"
logef(xi;yi)
1
KPK
j=1ef(xi;yj)#
fcould, for instance, be a bilinear function f(xi;yi) =zT
iWz0
i, wherezie(zjxi),z0
ie(zjyi),
andW2Rnnis a learned parameter matrix.
Bootstrapping. This is a simpliﬁed variant of contrastive learning. Given a related context xand
targety, 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 yfor different transformations of the
same input image.
Compression. A representation ze(jx)should contain enough information about the input x
to solve downstream tasks. Ideally, eshould also extract only the minimum amount of information
aboutxthat 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
minimizesI(X;ZjY).
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 V AE, 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 deﬁne 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 deﬁnition, the downstream tasks (which could include imitation,
classiﬁcation, 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
targetsxintozvectors 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 difﬁculty of the contrastive task scales with the number of negatives.
4
However, batches of the appropriate difﬁculty 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. Speciﬁcally, the target encoder’s weights tare updated
to track the main context encoder weights cusing the update rule t t+ (1)c.is
referred to as a momentum parameter, and is typically set to some value close to 1 (e.g. = 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 zoutput 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 V AEs 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 at. However,
for an encoder to be used for reinforcement learning or imitation, the representation must not depend
on the current action at. Thus, the encoder is only responsible for learning a zrepresentation of the
observation ot, 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 ﬁnd 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 brieﬂy 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, ﬁnger-spin, and reacher-easy tasks.
All three of these are popular benchmark tasks for deep RL and deep IL, and represent a range of
difﬁculties (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.
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.
5
Table 3: Design decisions for representation learning algorithms used in our experiments.
Algorithm Task Context Target Act Aug
Temporal CPC Contrastive otot+1 7 3
SimCLR Contrastive otot 7 3
V AE Reconstructive otot 7 7
Dynamics Reconstructive ot;atot+1 3 7
Inverse Dynamics Reconstructive ot;ot+1at 7 7
From MAGICAL we choose the MoveToRegion, MoveToCorner, and MatchRegions tasks, which
represent a range of difﬁculty levels (MoveToRegion being the easiest, and MatchRegions being
the hardest). For each task, MAGICAL deﬁnes 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 D=f(x0;a0);(x1;a1);:::gof observation–action tuples drawn
from a demonstrator, BC optimises the policy (ajx)to maximise the expected log likelihood,
L() = E
(x;a)D[log(ajx)]:
We combine BC with representation learning in two ways. First, we use RepL to pretrain all but the
ﬁnal layer of the policy, then ﬁne-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 (ajx)and a
discriminator D (x;a)that must distinguish ’s behaviour from that of the demonstrator. Using
alternating gradient descent, GAIL attempts to ﬁnd a and that attain the saddle point of
max
min

E
(x;a)[logD (x;a)]E
(x;a)D[log(1D (x;a))] +wHH()
:
HereHis an entropy penalty weighted by regularisation parameter wH0. 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 ﬁve repre-
sentation 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 standard deviation over at least
ﬁve 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 signiﬁcant 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.
6
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 signiﬁcant 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 V AE BC aug BC no aug
DMC cheetah-run 482 36 669 18 687 17 661 13 458 39 690 17 617 34
ﬁnger-spin 718 17 74817* 726 1 723 4 7516* 730 9 9404*
reacher-easy 774 24 89014 9079 89313 88020 874 21 452 34
Procgen coinrun-train 8.1 0.4 8.0 0.2 8.0 0.5 8.10.3 8.40.4 8.1 0.3 8.70.6*
fruitbot-train 3.2 1 16.2 1.2 17.5 1.9 15.4 1.5 17.5 1.5 18.3 1.9 11.4 0.6
jumper-train 8.1 0.2 8.0 0.4 7.9 0.6 7.5 0.6 7.9 0.6 8.1 1.2 7.1 1.2
miner-train 4.5 1.2 5.9 0.2 9.90.4 9.5 2.3 10.40.3* 9.8 0.3 8.1 0.3
coinrun-test 6.3 0.8 6.90.5 6.80.5 6.80.4 7.00.5 6.7 0.4 6.5 0.7
fruitbot-test -3 0.9 15.61.1 13.4 1.0 14.71.0 13.2 1.0 13.7 1.1 2.2 0.6
jumper-test 3.2 0.4 3.9 0.3 3.6 0.4 3.7 0.5 3.4 0.5 3.9 0.5 4.60.4
miner-test 0.6 0.1 2.6 0.1 2.6 0.4 3.10.4 2.70.3 2.7 0.4 0.8 0.1
MAGI- MatchRegions 0.42 0.04 0.42 0.04 0.42 0.03 0.41 0.01 0.42 0.03 0.43 0.02 0.28 0.08
CAL MoveToCorner 0.840.07 0.830.04 0.830.04* 0.800.02 0.780.06 0.78 0.05 0.72 0.04
MoveToRegion 0.820.02* 0.830.02* 0.820.01* 0.810.01* 0.810.05* 0.74 0.02 0.810.04*
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 signiﬁcant 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 V AE BC aug BC no aug
DMC cheetah-run 72314* 71623* 71711* 71616* 72412* 690 17 617 34
ﬁnger-spin 7556* 75512* 73215 725 12 7553* 730 9 9404*
reacher-easy 89819 90310* 88914 91218* 9038* 874 21 452 34
Proc- coinrun-train 8.0 0.4 7.1 0.3 8.0 0.5 8.60.5* 7.9 0.2 8.1 0.3 8.70.6*
gen fruitbot-train 17.0 0.7 6.6 1.4 13.4 1.9 11.4 0.7 15.4 1.0 18.3 1.9 11.4 0.6
jumper-train 7.9 0.5 8.1 0.4 8.0 0.4 8.0 0.3 8.30.5 8.1 1.2 7.1 1.2
miner-train 8.9 0.8 8.9 0.7 8.7 0.3 7.1 0.8 8.6 0.7 9.8 0.3 8.1 0.3
coinrun-test 6.4 0.4 6.0 0.5 6.6 0.3 6.2 0.5 6.90.4 6.7 0.4 6.5 0.7
fruitbot-test 10.9 0.7 3.3 1.1 8.5 1.5 6.4 1.2 10.4 1.6 13.7 1.1 2.2 0.6
jumper-test 3.4 0.3 4.80.2* 3.8 0.3 3.4 0.3 3.9 0.7 3.9 0.5 4.60.4*
miner-test 2.0 0.2 1.9 0.3 1.8 0.3 1.0 0.2 2.0 0.3 2.7 0.4 0.8 0.1
MAGI- MatchRegions 0.440.02 0.23 0.08 0.41 0.02 0.01 0.01 0.41 0.03 0.43 0.03 0.31 0.02
CAL MoveToCorner 0.78 0.07 0.30 0.22 0.76 0.05 0.02 0.02 0.820.06 0.80 0.05 0.70 0.09
MoveToRegion 0.760.02 0.35 0.24 0.74 0.01 0.47 0.07 0.770.02 0.75 0.02 0.780.04
BC pretraining results. In the pretraining setting, we see that none of our RepL algorithms consis-
tently 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 beneﬁt in some environments, but are sensitive to the choice of task.
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 (ﬁnger-spin,
coinrun-train, jumper-test, and MoveToRegion), adding augmentations leads to consistently worse
7
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 signiﬁcant 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 tCPC.
Env Task Dynamics InvDyn SimCLR tCPC V AE GAIL aug GAIL no aug
DMC cheetah-run 380 76 320 61 265 58 360 74 375 33 449 67 75 40
ﬁnger-spin 86814 8868* 800 23 748 72 86818 868 12 0 0
reacher-easy 53 24 73 51 21 23 118 88 122 89 221 162 89 88
Proc- coinrun-train 5.90.29* 5.850.51* 2.15 1.53 3.28 2.62 3.541.22 3.31 0.44 2.80 0.89
gen fruitbot-train -2.81 0.1-2.37 0.55 -2.47 0.15 -2.38 0.31 -2.49 0.22 -2.42 0.42 -2.63 0.30
jumper-train 3.31 0.31 3.17 0.40 3.36 0.53 2.69 1.31 3.450.70 3.44 0.52 3.470.53
miner-train 0.53 0.12 0.60 0.11 0.53 0.14 0.840.14* 0.51 0.07 0.65 0.10 0.770.18
coinrun-test 6.10.9* 5.910.16* 2.11 1.61 3.35 2.74 3.01 1.10 3.44 0.68 2.77 0.84
fruitbot-test -2.44 0.49 -2.65 0.24 -2.55 0.30 -2.65 0.14 -2.85 0.33 -2.44 0.50 -2.51 0.44
jumper-test 2.56 0.52 2.53 0.64 3.15 0.45 2.35 0.81 2.75 0.59 3.25 0.42 3.15 0.20
miner-test 0.36 0.04 0.57 0.07 0.55 0.24 0.870.15* 0.50 0.17 0.65 0.17 0.660.14
MAGI- MatchRegions 0.42 0.10 0.34 0.12 0.470.04 0.39 0.12 0.30 0.15 0.46 0.06 0.22 0.12
CAL MoveToCorner 0.48 0.09 0.45 0.10 0.520.07 0.550.15 0.620.11* 0.49 0.08 0.550.14
MoveToRegion 0.72 0.07 0.74 0.04 0.74 0.06 0.760.03 0.75 0.07 0.75 0.09 0.60 0.14
performance. This effect is particularly dramatic in ﬁnger-spin, which we believe is a result of the fact
that relevant objects in the environment always stay ﬁxed. Consequently, translational augmentations
don’t aid generalization, and rotational augmentations can be confused with true signal (since the
angle of the ﬁnger determines the ideal action). Because augmentation already yields large beneﬁts,
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
ﬁnding 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 signiﬁcant 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 ﬁnger-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].
4 Discussion & future work
Contrasting image classiﬁcation and imitation learning datasets. The use of self-supervised
representation learning for pretraining has met with notable success in image classiﬁcation [ 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.
STL-10
Labels ⟶
Images ⟶
MAGICAL Actions ⟶
States ⟶
Pr ocgen Actions ⟶
States ⟶
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 classiﬁcation than in IL tasks.The case of Behavioural
Cloning (BC) is particu-
larly illustrative. BC uses
the same optimization algo-
rithms, loss types, and net-
work architectures as other
forms of image classiﬁca-
tion, so if RepL is less help-
ful for BC than for other
forms of classiﬁcation then
it must be due to the choice
of training and evaluation
data. For the sake of illus-
8
coinrun-VAE-Acts
Acts
0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
coinrun-VAE-Returns
Rets
6
7
8
9
STL10-SimCLR-ResNet50Figure 3: t-SNE embedding of representations from a V AE 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).
trating differences in data
distributions, Fig. 2 compares the STL-10 dataset—a typical image classiﬁcation 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 yaxis.
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 inﬂuenced by ﬁne-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 classiﬁcation on STL-10 than for control on MAGICAL or Procgen.
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 ﬁne-grained
features necessary for action prediction. However, it could be that reward- and value-prediction
beneﬁt 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) ofﬂine Q-learning
does beneﬁt 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 V AE-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
9
Original
Saliency
Finger Spin Coinrun 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 inﬂuence over the network’s output.
Notice that SimCLR attends mainly to the foreground in DMC, and mainly to the background in
CoinRun.
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] beneﬁt more from RepL than algorithms that learn
policies.
The importance of using diverse benchmark tasks. Our experiment results in Table 5 showed
much greater beneﬁt 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 ﬁnger-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 ﬁnger-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 ﬁnger-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.
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-speciﬁc. Our analysis has identiﬁed several hypotheses that could help understand when
andwhere 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.
10
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
[1]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.
[2]Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple frame-
work for contrastive learning of visual representations. Proceedings of the 37th International
Conference on Machine Learning , 2020.
[3]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.
[4]Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas.
Reinforcement learning with augmented data. arXiv preprint arXiv:2004.14990 , 2020.
[5]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.
[6]Sam Toyer, Rohin Shah, Andrew Critch, and Stuart Russell. The MAGICAL benchmark for
robust imitation. arXiv preprint arXiv:2011.00401 , 2020.
[7]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.
[8]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.
[9]Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint
arXiv:1312.6114 , 2013.
[10] 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.
[11] 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 conﬁdence. In Proceedings of the Thirty-fourth Conference on Neural
Information Processing Systems , 2020.
[12] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive
predictive coding. arXiv preprint arXiv:1807.03748 , 2018.
[13] 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.
[14] 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 ,
2020.
[15] 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.
[16] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122 , 2018.
[17] 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.
11
[18] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised repre-
sentations for reinforcement learning. Proceedings of the 37th International Conference on
Machine Learning, Vienna, Austria, PMLR 119 , 2020. arXiv:2003.06417.
[19] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation
learning from reinforcement learning. arXiv preprint arXiv:2009.08319 , 2020.
[20] Jürgen Schmidhuber. Reinforcement learning in markovian and non-markovian environments.
InAdvances in Neural Information Processing Systems 3, NIPS’3 , pages 500–506, 1991.
[21] 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.
[22] Ian Fischer. The conditional entropy bottleneck. Entropy , 2020.
[23] Dean A Pomerleau. Efﬁcient training of artiﬁcial neural networks for autonomous navigation.
Neural Computation , 3(1):88–97, 1991.
[24] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in
neural information processing systems , pages 4565–4573, 2016.
[25] 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.
[26] Mengjiao Yang and Oﬁr Nachum. Representation matters: Ofﬂine pretraining for sequential
decision making. arXiv preprint arXiv:2102.05815 , 2021.
[27] Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse
reinforcement learning. arXiv preprint arXiv:1710.11248 , 2017.
[28] Siddharth Reddy, Anca D Dragan, and Sergey Levine. SQIL: Imitation learning via reinforce-
ment learning with sparse rewards. arXiv preprint arXiv:1905.11108 , 2019.
[29] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks:
Visualising image classiﬁcation models and saliency maps. arXiv preprint arXiv:1312.6034 ,
2013.
[30] Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep
variational reinforcement learning for pomdps. arXiv preprint arXiv:1806.02426 , 2018.
[31] 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.
[32] 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.
[33] Pei Yingjun and Hou Xinwen. Learning representations in reinforcement learning: An informa-
tion bottleneck approach. arXiv preprint arXiv:1911.05695 , 2019.
[34] 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.
[35] 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.
[36] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip
Bachman. Data-efﬁcient reinforcement learning with momentum predictive representations.
arXiv preprint arXiv:2007.05929 , 2020.
[37] 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.
[38] 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.
12
[39] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectiﬁers:
Surpassing human-level performance on ImageNet classiﬁcation. In Proceedings of the IEEE
Conference on Computer Vision and Pattern Recognition , pages 1026–1034, 2015.
[40] Ian Fischer and Alex Alemi. Ceb improves model robustness. Entropy , 2020.
13

Cynthia Chen
Sam Toyer
Cody Wild
Scott Emmons
Ian Fischer
Kuang-Huei Lee
Neel Alex
Steven Wang
Ping Luo
Stuart Russell
Pieter Abbeel
and Rohin Shah