Machine Learning Projects for Iterated
Distillation and Ampliﬁcation
Owain Evansa, William Saundersb, Andreas Stuhlmüllerc
July 3, 2019
Abstract
Iterated Distillation and Ampliﬁcation (IDA) is a framework for
training ML models. IDA is related to existing frameworks like
imitation learning and reinforcement learning, but it aims to solve
tasks for which humans cannot construct a suitable reward function
or solve directly.
This document reviews IDA and proposes three projects that
explore aspects of IDA. Project 1 applies IDA to problems in high-
school mathematics and investigates whether learning to decompose
problemscanimproveperformanceoversupervisedlearning. Project
2 applies IDA to neural program interpretation, where neural nets
are trained on the internal behavior (execution traces) of traditional
computer programs. Project 3 investigates whether adaptive com-
putation time (varying compute at inference time as a function of
the input) can improve the robustness and eﬃciency of IDA.
Our goal in outlining these projects is to generate discussion and
encourage research on IDA. We are not (as of June 2019) working
on these projects, but we are interested in collaboration.
Contents
0 Background on IDA 2
0.1 What is IDA? . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2
0.2 Examples . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4
0.3 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6
0.4 Project Goals . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7
1 Project 1: Amplifying Mathematical Reasoning 8
1.1 Motivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8
1.2 Project Directions . . . . . . . . . . . . . . . . . . . . . . . . . . 8
1.3 IDA for High School Mathematics . . . . . . . . . . . . . . . . . 10
1.4 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12
aUniversity of Oxford
bUniversity of Toronto
cOught. Correspondence to andreas@ought.org
1
2 Project 2: IDA for Neural Program Interpretation 12
2.1 Motivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12
2.2 Project Directions . . . . . . . . . . . . . . . . . . . . . . . . . . 13
2.3 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14
3 Project 3: Adaptive Computation 16
3.1 Motivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16
3.2 Project Directions . . . . . . . . . . . . . . . . . . . . . . . . . . 17
3.3 IDA, Fast and Slow . . . . . . . . . . . . . . . . . . . . . . . . . . 17
3.4 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18
0 Background on IDA
0.1 What is IDA?
Iterated Distillation and Ampliﬁcation (IDA) is a framework for training models
from data [1, 2]. IDA is related to and builds on existing frameworks like su-
pervised learning, imitation learning, and reinforcement learning. It is intended
for tasks where:
1. The goal is to outperform humans at the task or to solve hard instances.
2. It is not feasible to provide demonstrations or reward signals for super-
human performance at the task.1
3. Humans have some high-level understanding of the task and can also pro-
vide demonstrations or reward signals for easy instances of the task.
The idea behind IDA is to bootstrap using an approach similar to Alp-
haZero [3], but with a learned model of human reasoning steps taking the place
of the ﬁxed game simulator.
We will explain IDA in abstract terms and then describe concrete examples.
For broader discussion of IDA, including its relevance to the value alignment
problem, see [1, 2, 4, 5, 6, 7].
0.1.1 Technical description of IDA
We consider the following learning problem: We want to train a model (e.g. a
neural net) to solve tasks from the set T, where Tcontains a series of tasks that
get progressively harder for humans to solve. Formally, let T=SN
i=0Ti, where
tasks in Tnare harder than tasks in Tn1for all n.
We are given a training set which includes solutions to the easiest class of
problems T0and human demonstrations of decomposing tasks Tninto ﬁnitely
many slightly easier tasks in Tn1.
In IDA, the initial training steps are:
1More precisely: it is infeasible to provide large numbers of demonstrations or suﬃciently
dense reward signals for methods like imitation learning or RL to work well.
2
1.Mis trained by supervised learning2to reproduce the answers to the
easiest tasks T0.
2.Mis trained by supervised learning to imitate the human demonstrations
fordecomposing a task x2Tninto a set of tasks in Tn1and then
aggregating the solutions to solve x.
Steps (1) and (2) are analogous to the two parts of a recursive algorithm: the
base case and therecursive step . After initial training, Mcan solve tasks in T1
by ﬁrst decomposing them into tasks in T0.Mis then trained by supervised
learning on its ownsolutions to tasks in T1, enabling Mto solve T1tasksdirectly
(i.e. without decomposition). Solving a task in T1directly involves a single call
toM, while solving by decomposition into tasks in T0requires Mto be called
on each of the T0tasks.
This process of “training on its own solutions” can be iterated. Mis trained
by supervised learning to directly solve increasingly hard tasks in T, where
the target solutions (i.e. labels for supervised learning) are produced by M
itself via decomposition into tasks Mcan already solve. If supervised training
works perfectly at each iteration, then eventually Mcan solve any task in T
directly (with only single call to M). (This depends on the strong assumption
that humans can decompose all tasks Tninto tasks in Tn1— see [8, 11] for
discussion.)
We can now summarize IDA (see Figure 1). After training on steps (1) and
(2), the “base case” and “recursive step”, the following steps are repeated (for
n >0):
Ampliﬁcation Step
Msolves tasks in Tnby decomposing them into tasks in Tn1, which it
solves directly (without decomposition).
Distillation Step
Mis trained by supervised learning to solve tasks in Tndirectly, with
target solutions coming from the Ampliﬁcation Step.
It is called the “Ampliﬁcation Step” because it ampliﬁes the capability of model
M. While Mcan only solve tasks in Tn1directly, Mcan be used (via decom-
position and aggregation) to solve tasks in Tn.
In the “Distillation Step”, the slower ampliﬁed model (which makes multiple
calls to M) isdistilled into a faster process with a single call to M. This is
like distillation for neural nets [12], where a large net (or ensemble of nets) is
“distilled” into a smaller net that tries to capture the behavior of the large net.
In general, it is unlikely that distillation of the slower process will be perfect.
This can be addressed using RL-based distillation or by selectively choosing
when to use fast Mdirectly and when to fall back to the ampliﬁed model (see
Project 3).
2We describe IDA based on supervised learning, similar to [8], since this is what the
three projects in this document focus on. This can be substituted with RL or other training
schemes, see [9, 10]. We view answering questions, problem decomposition, and aggregation
of subproblem answers all as sequence-to-sequence problems, so we can train a single model
Mto solve them. We could also train distinct models.
3
x∈T0xa∈T0amplifydistillamplifyx∈T1xb∈T0x∈T1xa∈T1x∈T2xb∈T1distillx∈T2amplifyxa∈T2x∈T3xb∈T2Figure 1: Diagram showing the ﬁrst few Ampliﬁcation and Distillation steps in IDA
training. First (far left), Mis trained by supervised learning on each task x2T0.
Second, Msolveseachtask x2T1bydecomposingintotasks fxa; xbg 2T0andsolving
these directly. This eventually produces a dataset of solved tasks x2T1, which Mis
trained to solve directly by supervised learning (leftmost “distill” step). This process
is then repeated.
To simplify the exposition, we have presented Mas being trained to decompose
tasks and learn all base case solutions beforethe iterative training by Ampliﬁca-
tion and Distillation begins. In practice, these processes (learning to decompose
harder problems and gathering additional base case solutions from humans, and
learning to solve harder problems directly) would happen in parallel [8].
0.2 Examples
Having given an outline of IDA, we will describe two toy examples of solving
problems with IDA: integer multiplication and shortest path in graphs.3These
examples are intended to build intuition for IDA. They are not themselves prac-
tically relevant use cases for IDA.
Example 1: Multiplication
This toy example shows how one would use IDA to train a neural net Mto
multiply large integers. We assume that Mhas been pre-trained to add large
integers.
Following the pattern in the previous section, the initial training set contains
(1) simple multiplications ( base case ), and (2) demonstrations of decomposing
multiplications into simpler multiplications ( recursive step ).Mis trained as
follows:
1. Train Mon simple multiplications: single-digit multiplication and multi-
plication by 10. For example: 56 = 30,10234 = 2340 ,98 = 72.
2. Train Mtodecompose multiplications into simpler multiplications by
distribution and other algebraic rules. For example:
817 = 810 + 87
2234 = 2004 + 234 = (2010)4 + (204) + (34)
3The Multiplication example has not been implemented, but a version of the Shortest Path
example is implemented in [8].
4
After training Msuccessfully on (1) and (2), it would be possible to solve large
multiplications exactly by recursively decomposing down to the base case. This
is similar to how exact multiplication is computed by traditional algorithms or
by humans. A tree illustrating this is shown in Figure 2:
Solve 384 x 19300 x 1984 x 19300 x 10300 x 980 x 194 x 198×178×78×10384×19300×19300×10300×94×1980×1984×19
Figure 2: A multiplication problem could eventually be solved by decomposing all
the way down to simple multiplications (without any distillation). Only the ﬁrst two
levels of decomposition are shown here.
Recursing all the way to the base case would require many calls to the neural net
M. Instead, Mcould be trained incrementally using the IDA scheme described
above. If this was successful, Mcould eventually compute 38419or79332
2927in a single call, without the need to ever instantiate the fully expanded
slow process for large problems during training. This is a key property of IDA,
since recursive decompositions can (in general) result in a number of calls to
Mthat scales exponentially in the problem size, and so would be infeasible to
instantiate explicitly.
This example shows that IDA can easily fail if we use a standard neural
net as our model M. For example, a small MLP is not capable of learning to
multiply large integers in a single call. This is the case even for training the
MLP by supervised learning on ground-truth examples, e.g. training on pairs
((m; n); mn)form; n < 106. Training by IDA will generally make learning
morediﬃcult, because Mwill be trained by supervised learning on its own
answers to multiplication problems (instead of ground-truth answers).
This example also shows that IDA can fail regardless of the model M. It is
not generally possible to distill an exponential tree of calls to Minto a single
call to M. However, there are many AI problems where research aims to better
approximate an exponential-time computation. AlphaZero uses an IDA-like
algorithmtodistillanexponentialgame-treeexpansionintoafeedforwardneural
net. The neural net does not distill perfect play for Go or chess, but it achieves
impressive performance relative to humans (and formidable performance when
combined with MCTS).
Example 2: Shortest path in a graph
A recent paper by Christiano et al. [8] implements IDA and applies it to discrete
algorithms problems including union ﬁnd, wildcard search, and shortest path.
5
To apply IDA to ﬁnding the shortest path between nodes sandtin a directed
graph, we need an initial training set that covers the base case and recursive
step. These include:
1. A dataset of solutions to the easiest shortest path problems (for which
nodes sandtare adjacent).
2. Asetofdemonstrationsofdecomposingshortestpathproblemsintosmaller
shortest path problems and aggregating the results.
The decomposition in (2) is similar4to the following recursion, which is also
used in the dynamic programming algorithm for shortest path:
min-dist( s; t) = min (fmin-dist( s; x) + dist( x; t)jxadjacent to tg)
Here “ min-dist( s; t)” is the minimum path length between nodes sandt, and
“dist” is the distance between adjacent nodes. Christiano et al. use a Trans-
former model [13] as their model Mand they compare two ways of training the
Transformer to compute the shortest path. The ﬁrst approach is the IDA model
just described, which only gets labeled examples for the smallest shortest path
problems and must bootstrap to solve larger problems. The second approach is
regular supervised learning, where the model gets a large set of labeled examples
of all sizes. The main result is that the IDA approach is successful in getting
close to the performance of the supervised model. While this is a toy example,
it illustrates the general aim for IDA, which is training a model for tasks where
(a) there are only labels/demonstrations for easy problems, and (b) humans can
provide decompositions for going from harder to easier problems.
0.3 Related Work
As noted above, there are various discussions of IDA and its relevance to AI
alignment [1, 2, 6, 7]. These discussions are valuable as background but they
mostly abstract away the practical details of implementing an IDA system us-
ing ML. The only paper that actually implements IDA is the aforementioned
Christiano et al. However, the AlphaZero algorithm [3] is very similar to IDA [1]
and work applying AlphaZero (and related algorithms AlphaGo [14] and Expert
Iteration [15]) to chess/Go and to graph coloring [16] are relevant to thinking
about IDA experiments. IDA depends on bootstrapping and function approx-
imation, which are core topics in reinforcement learning [17, 18]. Recent work
on Deep Q-learning is especially relevant [19].
For IDA to help address value alignment problems for advanced ML sys-
tems, it would likely need to apply to tasks (and use decompositions) that
involve sophisticated reasoning in natural language. There is currently no pub-
lished research that is targeted at natural language tasks. Ought and OpenAI
have conducted preliminary experiments in an IDA-like setting with humans
in place of ML models. These experiments are aimed at shedding light on
whether diﬀerent incentive structures for IDA-like approaches (e.g. diﬀerent ob-
jective functions and reward signals) lead to aligned behavior. See examples
from Ought [20] and OpenAI [21] and stay in touch with Ought5for the latest
information on experiments.
4Christiano et al. use a slightly more complicated decomposition.
5https://ought.org
6
In the two toy examples above, IDA trains a model to imitate behavior.
This is an extension of supervised learning and imitation learning, and does not
involve reinforcement learning. However, as mentioned above, IDA is compati-
ble with RL [10, 9]. There is also a framework called “Debate”, which is closely
related to IDA and draws on self-play reinforcement learning as a method of
bootstrapping. Irving et al. [22] introduces Debate, explores analogies to com-
putational complexity theory, and describes links between Debate and IDA.
0.4 Project Goals
The rest of this document outlines three projects that could help clarify aspects
of IDA. These projects aim to extend Christiano et al [8] in a few diﬀerent ways:
The projects train a single model Mto perform a wide range of types
of decompositions. By contrast, Christiano et al. train a fresh model for
eachalgorithmicproblem(i.e.oneforeachofunionﬁnd,shortestpath,and
wildcard search). The decompositions for these algorithmic problems have
dynamic programming structure, so each model decomposes a problem
instance into smaller instances of the same kind of problem (e.g. a shortest
path problem is broken down into smaller shortest path problems).
The projects are well suited to investigating extrapolation and generaliza-
tion out of distribution. Projects 1 and 2 build on prior work in ML which
tests generalization performance in mathematics and neural programming
tasks. Project 3 is focused on general approaches (adaptive computation
time and calibration) to improving robust generalization.
In Christiano et al. the main goal is for a model trained by IDA to match
the performance of a model trained by supervised learning. The IDA
model has less labeled data and is trained by bootstrapping on its own
labels. While this is one possible goal for our projects, we also consider
the goal of improving test-time performance by making multiple calls to
the model using ampliﬁcation. Improving test-time performance by am-
pliﬁcation is similar to the way AlphaZero uses MCTS during competitive
matches to improve performance over the policy net.
There are many other possible projects on IDA. Research projects that push in
the following directions seem particularly valuable:
Produce theory and empirical knowledge about training IDA systems,
analogous to knowledge of how to train supervised learning or reinforce-
ment learning systems.
Connect IDA to existing work in ML. The projects below connect to lan-
guage modeling, neural programming, calibration for neural networks, and
adaptive computation time. IDA is also related to semi-supervised learn-
ing, deep reinforcement learning, dynamic programming, belief propaga-
tion, etc. A research project could explore and develop any of these con-
nections.
Solve problems with IDA that can’t be solved by other approaches. This
is the ultimate goal of IDA, and would draw interest from the larger ML
7
community. The projects below are not primarily aimed to achieve this,
but they may provide a useful ﬁrst step. Two ways in which IDA could
solve problems that can’t be solved by other approaches are:
–By distilling large (i.e. exponential in problem size) trees to a fast
machine learning model (as in AlphaZero).
–By learning decomposition steps from human data. This is for do-
mains (e.g. common-sense reasoning) where we are not able to write
down an algorithm to decompose problems.
1 Project 1: Amplifying Mathematical Reason-
ing
1.1 Motivation
Decomposition is a fundamental strategy for solving problems in mathematics.
Consider the following problem from high-school mathematics:
g(y) = y-2
f(x) = x\*g(x) + 3x^3
Find the derivative of f at x=1.
We can solve the problem by decomposing it into sub-problems which only
depend on parts of the whole problem. Here is one possible decomposition:
Sub-problem 1:
g(y) = y-2
f(x) = x\*g(x) + 3x^3
Write f as a polynomial in x in standard form.
Solution: f(x) = 3x^3 + x^2 - 2x
Sub-problem 2:
f(x) = 3x^3 + x^2 - 2x
Differentiate f.
Solution: f’(x) = 9x^2 + 2x - 2
Sub-problem 3:
f’(x) = 9x^2 + 2x - 2
Compute f’(1).
Solution: 9
These sub-problems could themselves be decomposed: the ﬁrst sub-problem de-
composes to substituting the function gand expanding the resulting expression.
1.2 Project Directions
The aim of this project is to train a model via IDA to solve mathematics prob-
lems. Following Christiano et al [8], we could use IDA at training time to
8
bootstrap from a small labeled dataset. This is like a student learning math by
solving problems from a textbook with no solutions at the back of the book.
Another possible goal is to use ampliﬁcation at test time to apply more com-
pute to harder problems, and so achieve better test-time performance than a
standard supervised model.
There are many possible choices of mathematics problem, including:
1. Mathematical problems in a formal language (e.g. theorem proving [23]).
2. High-school algebra (in a mix of natural language and math notation [24]).
3. So-called “word problems”, which are problems in natural language that
need to be translated into math. For example, many problems on the
American GMAT or GRE exams [25].
4. Advancedmathematicsproblems: competitionorOlympiadmath,university-
level proof-based math [26].
Problems in (1) do not require working with natural language. Many such
problems have a natural decomposition via brute-force search, similar to game-
tree search in Chess. These problems are a good testing ground for some aspects
of IDA and we encourage research on them. However this document focuses on
problems in (2)-(4) and especially (2). Problems in (2)-(4) are natural language
problems which don’t necessarily have an obvious decomposition strategy. This
makes them more similar to problems in areas outside mathematics such as
science, philosophy, and common-sense reasoning.
How can we tackle natural language mathematics problems in IDA? The
main question is how to produce decompositions. There are two basic options:
1. Have human experts produce decompositions of the problems.
2. WriteanalgorithmthatsolvesproblemsbydecompositionasinChristiano
et al [8]. This algorithm is likely to resemble classical AI approaches
(“GOFAI”).
Using human data is the more general approach, as it applies outside mathemat-
ics. However, it is unlikely that the usual way people solve problems would yield
the most useful decompositions. So, part of the research eﬀort is to work out
which kinds of decompositions help most and train human experts to produce
them.
For option (2) above, there are two kinds of algorithm for decomposition.
The ﬁrst kind is an eﬃcient algorithm that solves all problems in the class of
interest (as in the multiplication and shortest path examples). A neural net
trained using IDA is unlikely to perform better than such an algorithm. Ex-
periments with IDA would aim not at state-of-the-art performance but instead
at investigating certain aspects of IDA. We discuss this in the next section.
The second kind of algorithm solves problems in the class ineﬃciently and so
can only solve small problems in practice. In this case, it is possible that IDA
can achieve state-of-the-art performance by learning to distill decompositions
into a much more eﬃcient neural algorithm. However, for advanced mathemat-
ics problems in natural language, it is not clear what this ineﬃcient algorithm
9
would look like.6So there is a separate research project in investigating such
algorithms.
In the next section, we outline a project that uses high-school math prob-
lems in natural language. We suggest starting with algorithmically generated
decompositions and later extending this to decompositions provided by humans.
1.3 IDA for High School Mathematics
1.3.1 Task and Dataset
Saxton et al. [24] introduce a dataset of high-school level mathematics problems
in natural language. The problems cover arithmetic, algebra, diﬀerentiation,
probability, and number theory. Here are some examples:
Question: Solve -42\*r + 27\*c = -1167 and 130\*r + 4\*c = 372 for r.
Answer: 4
Question: Simplify sqrt(200)\*2 + sqrt(200) + sqrt(200) + -4.
Answer: -4 + 40\*sqrt(2)
Question: Let u(n) = -n^3 - n^2. Let e(c) = -2\*c^3 + c. Let
f(j) = -118\*e(j) + 54\*u(j). What is the derivative of f(a)?
Answer: 546\*a^2 - 108\*a - 118
Question: Three letters picked without replacement from
qqqkkklkqkkk. Give prob of sequence qql.
Answer: 1/110
Question: What are the prime factors of 235232673?
Answer: 3, 13, 19, 317453
Question: Let j = -5 - 28. Is j/6\*(-14) a composite number?
Answer: True
The problems are generated by an algorithm, so there is unlimited training data.
The dataset includes test sets for both interpolation andextrapolation . The ex-
trapolation questions have quantities that vary outside of the range encountered
on the training set.
The paper includes baselines for models trained by supervised learning. The
questions and answers are both represented as strings, so it can be treated as a
sequence-to-sequence problem. The best performing model is a standard Trans-
former with 30M parameters, which achieves 76% model accuracy (probability
of a correct answer) on interpolation and 50% on extrapolation.
1.3.2 Approach with IDA
How should one generate decompositions for training IDA to solve these math
problems? Oneoptionistowriteanalgorithmfordecomposingproblems, rather
than collecting human decompositions. The problems were generated using a
compositional algorithm: consider the last example in the list above, which
6If mathematics problems are fully formalized, we can search over formal proofs. But if
the mathematics is informal, it is much harder to provide an algorithm that would eventually
(given arbitrary amounts of time and compute) solve the problems.
10
combines algebra (specifying an integer using equations) and number theory
(checking if the integer is composite). This algorithm can be run in “reverse” to
help generate decompositions. However, it is not obvious how close the resulting
decompositions would be to decompositions that are a good ﬁt for IDA training.
(For instance, the best decompositions could be more or less ﬁne-grained.)
One objective would be to do better than the supervised baseline at test time
by applying ampliﬁcation (i.e. decomposing the problem and using multiple
calls to the model M). In particular, ampliﬁcation promises to do better at
extrapolation to bigger problems or problems with larger numerical quantities.
Another objective is to train from a smaller number of labels using distillation
and try to rival the performance of the supervised baseline.
1.3.3 Non-Ampliﬁcation Baselines
The simplest baseline is supervised learning (as in Saxton et al [24]). The math
problems and solutions are represented as strings, and the model is trained to
mapstringstostrings. Anotherbaselinewouldmakeuseofthesamedecomposi-
tion training data as IDA. Instead of training the model to decompose problems,
we could use the decomposition as an auxiliary objective. The model would be
trained to produce both the decomposition and the answer. An alternative ap-
proach is to train a model to take strings as input and then produce output
suitable to be fed into a symbolic math system (e.g. SymPy).
1.3.4 Questions to investigate
The project would seek to investigate some of the following questions:
What is the test-time performance of applying ampliﬁcation vs. a baseline
that makes a single call to the model (distillation) vs. a baseline that was
trained by supervised learning?
What is the performance on extrapolation and on oﬀ-distribution prob-
lems?
What is the performance of a distilled model trained by IDA (from a
small number of labeled examples) vs. the supervised baseline (with a
large labeled training set)?
How does performance vary with diﬀerent kinds of decomposition strate-
gies?
How robust is the ampliﬁed model to noise in the training data and to
approximation error in the neural net?
Can we ﬁnd a decomposition strategy that makes the model more robust
to errors/noise?
When training a model to solve sub-problems, we need to pick some dis-
tribution over sub-problem examples. How does this impact performance?
Are there principles for generating training data in the most useful way
for IDA?
11
Using ampliﬁcation at test time is one way to vary the amount of compute
used to solve problems. Another approach is to use recurrent models and
adaptive computation time. How does this compare to IDA? (See Project
3 for more discussion and references).
Can we train the model from decompositions provided by humans? What
kind of decompositions should we use? Can these be augmented by algo-
rithmically generated decompositions?
1.4 Related Work
Analysing Mathematical Reasoning Abilities of Neural Models
(Saxton et al.) [24]
Dataset: https://github.com/deepmind/mathematics\_dataset .
The paper by Saxton et al. discussed above.
Neural Arithmetic Logic Units (Trask et al.) [27]
A specialized neural net architecture for doing arithmetic. It achieved
state-of-the-art results (as of 2018) on various tasks, including pure arith-
metic and arithmetic combined with vision/natural language.
Program Induction by Rationale Generation (Ling et al.) [25]
Dataset: https://github.com/deepmind/AQuA
This paper introduces a dataset of mathematical word problems (based
on US standardized tests). Humans had to “show their work” while solv-
ing the problems. The paper has a model that learns to generate these
intermediate steps (in addition to learning to solving the problems).
Sigma Dolphin
microsoft.com/en-us/research/project/sigmadolphin-2/
Dataset of natural language math problems, taken from Yahoo Answers.
HOList: An Environment forMachine Learningof Higher-Order
Theorem Proving (Bansal et al.) [23].
This dataset contains fully formalized proofs for a large number of theo-
rems and a framework for training ML systems to produce proofs.
2 Project 2: IDA for Neural Program Interpre-
tation
2.1 Motivation
Many algorithms (e.g. matrix arithmetic, shortest path, sorting) involve de-
composing problems into progressively smaller problems and then aggregating
results. More generally, computer programs in high-level languages decompose
tasks into progressively smaller tasks, and ultimately into the primitive oper-
ations of the language. We can explore the capabilities of IDA by training a
model on the decompositions used in the execution of computer programs.
12
Training IDA on decompositions from programs is closely related to research
onNeuralProgramInterpretation(seeReedanddeFreitas[28]andrelatedwork
below) or “NPI”. In NPI, a model is trained to mimic the internal behavior of an
algorithm and not just its input-output behavior. Moreover, the model trains
not just on the primitive operations of the algorithm but on its hierarchical de-
composition (i.e. the way procedures call other procedures). As with IDA, one
motivation for learning this internal behavior is to achieve stronger generaliza-
tion. Another motivation is to integrate hierarchical discrete computation with
the sort of pattern recognition in high-dimensional spaces enabled by machine
learning (see third experiment in [28]).
2.2 Project Directions
The project could focus on either of the following areas:
1.Distillation of programs
In contrast to most work on NPI, the goal of IDA experiments could be
to use program decompositions as training data to learn more eﬃcient
“neural” programs. The idea is to distill elaborate computations into a
single call to a neural net, or to combine the exact (slow) computation
with distillation (as in AlphaZero).
2.NPI within a more general framework for learning from decom-
positions
Much of the NPI work uses environments and architectures designed spe-
cially for NPI. For IDA, we would aim to replicate NPI results, but in a
framework that would also allow learning from human decompositions in
natural language.
We expect distillation of programs to be challenging. If we take a complicated
algorithm and try to distill it into a neural net, there is likely to be approxima-
tion error—and errors will usually be larger for inputs that are oﬀ the training
distribution. This will cause problems both during IDA’s iterative training pro-
cedure and also at test time. Part of the project would be to investigate how
well distillation works for diﬀerent kinds of programs and for diﬀerent ways of
organizing the training curriculum. If the distilled model was calibrated , then
IDA could recognize when distillation was likely to fail and fall back on us-
ing ampliﬁcation. Project 3 (below) explores this “adaptive computation” or
“meta-reasoning” approach.
2.2.1 Decision Points
There are many choices to make in devising IDA experiments in the NPI setting:
Which programs should we try to learn?
The research on NPI has focused on basic algorithms for tasks like integer
arithmetic and sorting. Applying IDA to these basic algorithms could be a
good starting point, as they allow comparison to existing work. However,
it isn’t clear how much experience with these algorithms would generalize
to other applications of IDA. It could be good to consider algorithms that
work with databases or knowledge bases, or to consider algorithms that
13
operate on human-readable structures like natural languages or images.
We also think that pure functional programming is a better ﬁt for IDA
than imperative programming. See [11, 29] for relevant discussion.
What kind of built-in operations and environments should we use?
In existing work on NPI, the neural net is given outputs that correspond
to basic operations on data. This makes it easier to learn algorithms
that depend on those basic operations. For IDA, it would be ideal to
learn these operations from examples. (If we were learning from human
decompositions, wemightnotknowaboutthese“basicoperationsondata”
ahead of time).
What kind of performance objective should we focus on?
Some work on NPI has focused on getting perfect performance on narrow
algorithmic tasks. It’s not clear if this is the right objective for IDA. We
might care about (a) generalizing well to much larger inputs mostof the
time (but not in the worst case), (b) being robust to distribution shift,
and (c) having one neural net learn a wide variety of algorithms.
2.2.2 Non-Ampliﬁcation Baselines
For learning to predict the next step of a program from examples, simpler ML
methods (random forests, logistic regression, etc.) may perform better than
neural networks. For performing the same task as a traditional program, any
neural program interpreter working with polynomial-sized trees will add signif-
icant overhead and so is unlikely to \*improve\* performance. One could also
consider tasks that require ML to process the inputs (e.g. tasks involving im-
ages). The baseline in this case would be a program that defers some decisions
to a classiﬁer.
2.3 Related Work
Neural Programmer-Interpreters (Reed and de Freitas) [28]
A quote from the introduction about the motivation for the paper:
“We may envision two approaches to providing supervision. In one,
we provide a very large number of labeled examples, as in object
recognition, speech and machine translation. In the other, the ap-
proach followed in this paper, the aim is to provide far fewer labeled
examples, but where the labels contain richer information allowing
the model to learn compositional structure. While unsupervised and
reinforcement learning play important roles in perception and motor
control, other cognitive abilities are possible thanks to rich supervi-
sion and curriculum learning. This is indeed the reason for sending
our children to school.”
Summary of the approach:
The model (called the “neural programmer-interpreter”) has a single in-
ference core for executing three diﬀerent programs (addition, sorting, ro-
tating CAD models). So one set of LSTM parameters are used to execute
14
all programs. However, the diﬀerent programs are stored as diﬀerent em-
beddings, stored in a learnable persistent memory.
The sequence of the model’s actions depends on the environment state
and action history.
For each program (addition, sorting, rotating CAD models), the model is
given a speciﬁc environment and set of actions in that environment. The
model is trained to compose these actions. For integer addition, there
are 1-D arrays and read-only pointers (for reading inputs), as well as 2-D
scratchpads and output arrays. For rotating CAD models, there’s a CAD
renderer with controllable elevation and azimuth movements.
The LSTM input of the previous computation step is a vector embedding,
rather than text.
Making Neural Programming Architectures Generalize via Recursion
(Cai et al.) [30]
Builds on Reed and de Freitas (above) and has no new machinery. They simply
allow a function to call itself. This means it can solve instances of arbitrary size
using recursive function calls, each of which have bounded length. This allows
for generalization oﬀ the training distribution. The hidden state of the LSTM
controller is reset (to zero) at each subprogram call, but the environment state is
not reset. They learn the recursion termination condition. They achieve 100%
generalization on all tasks (albeit for simple tasks).
Parametrized Hierarchical Procedures for Neural Programming (Fox
et al.) [31]
Their model is related to IDA for neural programming (it learns PHPs that
can recursively call other PHPs). Their tasks are limited to addition and to a
building a tower in gridworld. They provide a “weak supervision” motivation:
learn from a mix of traces that show what information should be remembered
from previous states and also from traces that omit that information.
Neural Program Lattices (Li et al.) [32]
This paper has a “weak supervision” motivation: learn from mix of full traces
and traces without program calls/arguments. They generalizes to 500-digit
addition, but not to 1000-digit addition.
Improving the Universality and Learnability of Neural Programmer-
Interpreters with Combinator Abstraction (Xiao et al.) [33]
Adds combinator abstraction from functional programming. One motivation
is to use reinforcement learning to learn programs without supervision. The
combinator makes the search space smaller.
15
Adaptive Neural Compilation (Bunel et al.) [34]
Takes programs, translates them into a diﬀerentiable form, then uses backprop-
agation to optimize them. They optimize programs to be more eﬃcient on a
restricted training distribution of problems.
Recent Advances in Neural Program Synthesis (Kant) [35]
This paper provides a summary of diﬀerent approaches to NPI and Neural
Program Synthesis.
LearningCompositionalNeuralProgramswithRecursiveTreeSearch
and Planning (Pierrot et al.) [36]
NPIapproachthat“incorporatesthestrengthsofNeuralProgrammer-Interpreters
(NPI) and AlphaZero”. While relevant to IDA, this approach diﬀers in impor-
tant ways.
3 Project 3: Adaptive Computation
3.1 Motivation
An ML model exhibits “adaptive computation” if it intelligently varies its com-
putations for diﬀerent inputs. For example:
1. The model selects which type of computation to run: e.g. between a slow
tree search, a large neural net, and a small neural net.
2. The model prioritizes possible computations: e.g. which node to expand
next in a tree search.
3. The model determines how long to run a ﬁxed computation: e.g. how
many MCTS samples, how many steps to run an RNN, etc.
A principled way to adapt computations is by “meta-level control” or “meta-
reasoning”. Meta-level control means applying ideas from optimal control and
Bayesian decision theory to selecting computations. The idea is to treat the
choice of computations as just another learning and planning problem. The
choice of computations can be optimized using end-to-end supervised learn-
ing [37], model-free reinforcement learning [38], or Bayesian decision theory [39,
40, 41].
Adaptive computation and meta-level control are not a major focus in cur-
rent deep learning research. Yet human cognition is often adaptive: people
dynamically decide how much time to spend on a task based on its perceived
diﬃculty. This is important when humans try to develop new ideas, which can
require anything from hours to years of thinking [42, 43]. By contrast, many
applications of ML have the following proﬁle:
1.Training time: Large amounts of compute and time are permitted. For
example, training might take months and use many CPUs and GPUs.
16
2.Test/deployment time: There are very strong constraints on compute.
For example, the trained model might be part of a web application and
so must perform inference in a fraction of a second.
Not all ML algorithms have this proﬁle. When AlphaZero [3] plays competitive
matches, it has time to make many calls to a neural net for each move. Alp-
haZerousesMCTStoinvestigatepromisingmovesfromthecurrentposition. As
ML is applied to more tasks for which humans spend a long time thinking (e.g.
mathematics, science, business strategy), the proﬁle of AlphaZero may become
more common.
IDA is well suited to adaptive computation. Training by IDA produces the
following algorithms:
A quick but potentially inaccurate algorithm for solving problems, pro-
duced by distillation. This corresponds to a single call to the model M
(using the terminology from Section 0.1).
An “anytime” algorithm for solving the same problems, which produces
more accurate or reliable answers as a function of doing more compute.
This is obtained by decomposing the problem using the learned decom-
position strategy (“Ampliﬁcation”) and involves making multiple calls to
M.
3.2 Project Directions
AdaptivecomputationforIDAcanbeappliedeitherduringtheiterativetraining
scheme or at test/deployment time:
1.Training time
The goal during training is to bootstrap by using ampliﬁcation to solve
progressively harder problems. Adaptive computation time could be ap-
plied to select how much computation to apply to a given problem in the
training set. For example, if the distilled model already performs perfectly
on some class of problems, there is no need to run the slow ampliﬁed pro-
cess (with many calls to model M) on this class. It’s also possible to
apply active learning – automatically selecting problems that are more
informative about the objective.
2.Test time
The goal is to perform well on a set of test problems given a time bud-
get. If the test set is received in batch, then the algorithm would ideally
spend more of its time budget on the harder instances. This requires the
algorithm to recognize harder instances that will beneﬁt from more com-
pute time (via ampliﬁcation). There are also more ﬁne-grained questions
about how to use the compute budget. For example, when using ampli-
ﬁcation, which problems should be decomposed ﬁrst and how far should
the recursive decomposition go?
3.3 IDA, Fast and Slow
There are many ways to use adaptive computation as part of IDA. As a starting
point, we describe a simple form of adaptive computation, where the algorithm
17
decidesbetweena fast(distilled)processforsolvingproblemsanda slowprocess
that uses a ﬁxed amount of ampliﬁcation. The fast process makes a single call
to a neural net, while the slow process makes at most ncalls to the same net.
For each input problem, the algorithm runs the fast process and has to decide
whetherto alsorunthe slowprocess. Thegoalis tobalance thegreateraccuracy
of the slow process with a time/compute cost that is speciﬁed as part of the
problem statement.
This decision of whether to run the slow process could be trained end-to-
end, by always running both processes during training (as in [37]). A diﬀerent
approach would be to use a calibrated model for the fast process and decide
whethertocalltheslowprocessbasedontheconﬁdenceofthiscalibratedmodel.
Which tasks would be a good testing ground for adaptive computation? The
mathematics and neural programming tasks (Projects 1 and 2 above) are useful
for testing purposes, because it’s easy to vary the amount of computation that
is necessary and suﬃcient for solving a problem. It’s also valuable to explore
tasks which have an “anytime” structure, where approximate solutions to the
task can be improved with additional compute. This would include planning,
natural language reasoning, chess, and strategic videogames.
How would this project diﬀer from Projects 1 and 2? The aim in those
projects is to explore the performance of IDA on challenging tasks. It’s not
clear that any clever adaptive computation strategy is required to do well on
these tasks (though we can’t rule it out). In a related example, AlphaZero did
not select the number of MCTS samples as a function of the board position,
but instead used a ﬁxed proportion of the remaining match time. Project 3, on
the other hand, is all about investigating adaptive computation for IDA. The
motivation is that adaptive computation is likely to play an important role as
IDA is applied to increasingly challenging tasks.
3.3.1 Non-Ampliﬁcation Baselines
Foranyadaptivecomputationapproach, itisimportanttocompareagainstnon-
adaptive approaches (which use a ﬁxed amount of compute) and also against
simple heuristics for scaling up compute for harder instances. Ampliﬁcation
could also be compared to adaptive approaches that do not use ampliﬁcation
(e.g. something like [37]).
3.4 Related Work
Attending to Mathematical Language with Transformers (Wangper-
awong) [44]
Thereisadataset(availableat https://github.com/tensorflow/tensor2tensor )
and paper for addition/multiplication/subtraction of numbers. The paper ex-
plores Transformers that can use more compute for larger instances.
AdaptiveComputationTimeforRecurrentNeuralNetworks(Graves)[37]
Adaptive Computation Time for RNNs adds the probability of halting compu-
tation at current step to the output of RNN. The idea is to train the adaptive
18
RNN end-to-end: running long computations at training time, which allows
diﬀerentiation of the halting probability.
Comparing Fixed and Adaptive Computation Time for Recurrent
Neural Networks (Fojo et al.) [45]
This paper presents experiments claiming to show that adaptive computation
time for RNNs (as in [37] above) is not needed because similar performance
is achieved with a regular RNN (which takes a ﬁxed number of steps between
predictions).
Universal Transformers (Dehghani et al.) [46]
ThepaperappliesAdaptiveComputationTimetotheTransformerarchitecture.
ItshowsimprovedperformanceonbAbItaskswithAdaptiveComputationTime
On Calibration of Modern Neural Networks (Guo et al.) [47]
This paper tests the calibration of convolutional neural networks and ﬁnds that
they are poorly calibrated by default. They show that temperature scaling
(learning a temperature parameter that scales the softmax inputs) does a good
job of getting the network to be calibrated. They indicate a tradeoﬀ between
accuracy and calibration (see their Figure 3).
Simple and Scalable Predictive Uncertainty Estimation using Deep
Ensembles (Lakshminarayanan et al.) [48]
They get calibrated uncertainty estimates for deep neural nets by:
Using a proper scoring rule for the loss
Performing adversarial training
Using an ensemble of networks
Principles of Metalevel Control (Hay) [39]
Ph.D. thesis on metalevel control and metareasoning. Includes an excellent
introduction to the subject and a review of previous work. The technical con-
tributions include applying Bayesian decision theory to Bandit problems and
applying RL to tree-search (i.e. learning a tree-search policy rather than just
building in MCTS).
Learning to Search with MCTSnets (Guez et al.) [49]
They show how to learn the hyperparameters of Monte-Carlo Tree search end-
to-end, by playing single-player gridworld game Sokoban. They learn a more
sample eﬃcient search than regular MCTS.
19
Acknowledgements
We thank Roger Grosse, Paul Christiano, Ondrej Bajgar and David Abel for
valuable discussions. This work was supported by a grant from the Future of
Life Institute (RFP2-178).
References
[1] Paul Christiano and Ajeya Cotra. Iterated distilla-
tion and ampliﬁcation. https://ai-alignment.com/
iterated-distillation-and-amplification-157debfd1616 . Accessed:
2019-05-17.
[2] Paul Christiano and Dario Amodei. Learning complex goals with iterated
ampliﬁcation. https://openai.com/blog/amplifying-ai-training/ .
Accessed: 2019-05-17.
[3] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou,
Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Ku-
maran, Thore Graepel, et al. A general reinforcement learning algo-
rithm that masters chess, shogi, and go through self-play. Science,
362(6419):1140–1144, 2018.
[4] Stuart Russell, Daniel Dewey, and Max Tegmark. Research priorities for
robust and beneﬁcial artiﬁcial intelligence. Ai Magazine , 36(4):105–114,
2015.
[5] Paul Christiano. Prosaic ai alignment. https://ai-alignment.com/
prosaic-ai-control-b959644d79c2 . Accessed: 2019-05-17.
[6] Paul Christiano. Towards formalizing uni-
versality. https://ai-alignment.com/
towards-formalizing-universality-409ab893a456 . Accessed: 2019-
05-17.
[7] Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini, and
Shane Legg. Scalable agent alignment via reward modeling: a research
direction. arXiv preprint arXiv:1811.07871 , 2018.
[8] Paul Christiano, Buck Shlegeris, and Dario Amodei. Supervising strong
learners by amplifying weak experts. arXiv preprint arXiv:1810.08575 ,
2018.
[9] Paul Christiano. Benign model-free rl. https://ai-alignment.com/
benign-model-free-rl-4aae8c97e385 . Accessed: 2019-05-17.
[10] William Saunders. Reinforcement learning in
the iterated ampliﬁcation framework. https://
www.alignmentforum.org/posts/fq7Ehb2oWwXtZic8S/
reinforcement-learning-in-the-iterated-amplification . Accessed:
2019-05-17.
20
[11] Andreas Stuhlmüller. Factored cognition. https://ought.org/
presentations/factored-cognition-2018-05 . Accessed: 2019-05-17.
[12] Geoﬀrey Hinton, Oriol Vinyals, and Jeﬀ Dean. Distilling the knowledge in
a neural network. arXiv preprint arXiv:1503.02531 , 2015.
[13] AshishVaswani, NoamShazeer, NikiParmar, JakobUszkoreit, LlionJones,
Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you
need. In Advances in neural information processing systems , pages 5998–
6008, 2017.
[14] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre,
George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda
Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep
neural networks and tree search. nature, 529(7587):484, 2016.
[15] Thomas Anthony, Zheng Tian, and David Barber. Thinking fast and slow
with deep learning and tree search. In Advances in Neural Information
Processing Systems , pages 5360–5370, 2017.
[16] Jiayi Huang, Mostofa Patwary, and Gregory Diamos. Coloring big graphs
with alphagozero. arXiv preprint arXiv:1902.10162 , 2019.
[17] Dimitri P Bertsekas and John N Tsitsiklis. Neuro-dynamic programming ,
volume 5. Athena Scientiﬁc Belmont, MA, 1996.
[18] Sutton Richard, Barto, and G Andrew. Reinforcement learning: an Intro-
duction. MIT Press, 2018.
[19] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel
Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K
Fidjeland, Georg Ostrovski, et al. Human-level control through deep rein-
forcement learning. Nature, 518(7540):529, 2015.
[20] Andreas Stuhlmüller. Ought 2018-12-31 progress update. https://ought.
org/updates/2018-12-31-progress-update . Accessed: 2019-05-17.
[21] GeoﬀreyIrvingandAmandaAskell. Aisafetyneedssocialscientists. Distill,
2019. https://distill.pub/2019/safety-needs-social-scientists.
[22] Geoﬀrey Irving, Paul Christiano, and Dario Amodei. Ai safety via debate.
arXiv preprint arXiv:1805.00899 , 2018.
[23] Kshitij Bansal, Sarah M Loos, Markus N Rabe, Christian Szegedy,
and Stewart Wilcox. Holist: An environment for machine learning
of higher-order theorem proving (extended version). arXiv preprint
arXiv:1904.03241 , 2019.
[24] David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli.
Analysingmathematicalreasoningabilitiesofneuralmodels. arXiv preprint
arXiv:1904.01557 , 2019.
[25] Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program
induction by rationale generation: Learning to solve and explain algebraic
word problems. arXiv preprint arXiv:1705.04146 , 2017.
21
[26] Mohan Ganesalingam and William Timothy Gowers. A fully automatic
theorem prover with human-style output. Journal of Automated Reasoning ,
58(2):253–291, 2017.
[27] Andrew Trask, Felix Hill, Scott E Reed, Jack Rae, Chris Dyer, and Phil
Blunsom. Neural arithmetic logic units. In Advances in Neural Information
Processing Systems , pages 8035–8044, 2018.
[28] Scott Reed and Nando De Freitas. Neural programmer-interpreters. arXiv
preprint arXiv:1511.06279 , 2015.
[29] Paul Christiano. Meta-execution. https://ai-alignment.com/
meta-execution-27ba9b34d377 . Accessed: 2019-05-17.
[30] Jonathon Cai, Richard Shin, and Dawn Song. Making neural programming
architectures generalize via recursion. arXiv preprint arXiv:1704.06611 ,
2017.
[31] Roy Fox, Richard Shin, Sanjay Krishnan, Ken Goldberg, Dawn Song, and
Ion Stoica. Parametrized hierarchical procedures for neural programming.
ICLR 2018 , 2018.
[32] Chengtao Li, Daniel Tarlow, Alexander L Gaunt, Marc Brockschmidt, and
Nate Kushman. Neural program lattices. 2016.
[33] Da Xiao, Jo-Yu Liao, and Xingyuan Yuan. Improving the universality and
learnability of neural programmer-interpreters with combinator abstrac-
tion.arXiv preprint arXiv:1802.02696 , 2018.
[34] Rudy R Bunel, Alban Desmaison, Pawan K Mudigonda, Pushmeet Kohli,
and Philip Torr. Adaptive neural compilation. In Advances in Neural
Information Processing Systems , pages 1444–1452, 2016.
[35] Neel Kant. Recent advances in neural program synthesis. arXiv preprint
arXiv:1802.02353 , 2018.
[36] Thomas Pierrot, Guillaume Ligner, Scott Reed, Olivier Sigaud, Nicolas
Perrin, Alexandre Laterre, David Kas, Karim Beguir, and Nando de Fre-
itas. Learning compositional neural programs with recursive tree search
and planning. arXiv preprint arXiv:1905.12941 , 2019.
[37] Alex Graves. Adaptive computation time for recurrent neural networks.
arXiv preprint arXiv:1603.08983 , 2016.
[38] Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and
Pieter Abbeel. RL2: Fast reinforcement learning via slow reinforcement
learning. arXiv preprint arXiv:1611.02779 , 2016.
[39] Nicholas James Hay. Principles of Metalevel Control . PhD thesis, UC
Berkeley, 2016.
[40] Sebastian Schulze and Owain Evans. Active reinforcement learning with
monte-carlo tree search. arXiv preprint arXiv:1803.04926 , 2018.
22
[41] Arthur Guez, David Silver, and Peter Dayan. Eﬃcient bayes-adaptive re-
inforcement learning using sample-based search. In Advances in neural
information processing systems , pages 1025–1033, 2012.
[42] Daniel Kahneman. Thinking, fast and slow . Macmillan, 2011.
[43] Owain Evans, Andreas Stuhlmüller, Chris Cundy, Ryan Carey, Zachary
Kenton, Thomas McGrath, and Andrew Schreiber. Predicting human de-
liberative judgments with machine learning. Technical report, Technical
report, University of Oxford, 2018.
[44] ArtitWangperawong. Attendingtomathematicallanguagewithtransform-
ers.arXiv preprint arXiv:1812.02825 , 2018.
[45] Daniel Fojo, Víctor Campos, and Xavier Giró-i Nieto. Comparing ﬁxed and
adaptive computation time for recurrent neural networks. arXiv preprint
arXiv:1803.08165 , 2018.
[46] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and
Łukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819 ,
2018.
[47] Chuan Guo, Geoﬀ Pleiss, Yu Sun, and Kilian Q Weinberger. On calibra-
tion of modern neural networks. In Proceedings of the 34th International
Conference on Machine Learning-Volume 70 , pages 1321–1330. JMLR. org,
2017.
[48] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Sim-
ple and scalable predictive uncertainty estimation using deep ensembles.
InAdvances in Neural Information Processing Systems , pages 6402–6413,
2017.
[49] Arthur Guez, Théophane Weber, Ioannis Antonoglou, Karen Simonyan,
Oriol Vinyals, Daan Wierstra, Rémi Munos, and David Silver. Learning to
search with mctsnets. arXiv preprint arXiv:1802.04697 , 2018.
23

Owain Evans
William Saunders
Andreas Stuhlmüller


The MAGICAL Benchmark for Robust Imitation
Sam Toyer Rohin Shah Andrew Critch Stuart Russell
Department of Electrical Engineering and Computer Sciences
University of California, Berkeley
{sdt,rohinmshah,critch,russell}@berkeley.edu
Abstract
Imitation Learning (IL) algorithms are typically evaluated in the same environment
that was used to create demonstrations. This rewards precise reproduction of
demonstrations in one particular environment, but provides little information about
how robustly an algorithm can generalise the demonstrator’s intent to substantially
different deployment settings. This paper presents the MAGICAL benchmark suite,
which permitssystematicevaluation of generalisation by quantifying robustness
to different kinds of distribution shift that an IL algorithm is likely to encounter
in practice. Using the MAGICAL suite, we conﬁrm that existing IL algorithms
overﬁt signiﬁcantly to the context in which demonstrations are provided. We also
show that standard methods for reducing overﬁtting are effective at creating narrow
perceptual invariances, but are not sufﬁcient to enable transfer to contexts that
require substantially different behaviour, which suggests that new approaches will
be needed in order to robustly generalise demonstrator intent. Code and data for
the MAGICAL suite is available athttps://github.com/qxcv/magical/.
1 Introduction
Imitation Learning (IL) is a practical and accessible way of programming robots to perform useful
tasks [ 6]. For instance, the owner of a new domestic robot might spend a few hours using tele-
operation to complete various tasks around the home: doing laundry, watering the garden, feeding
their pet salamander, and so on. The robot could learn from these demonstrations to complete the
tasks autonomously. For IL algorithms to be useful, however, they must be able to learn how to
perform tasks from few demonstrations. A domestic robot wouldn’t be very helpful if it required
thirty demonstrations before itﬁgured out that you are deliberately washing your purple cravat
separately from your white breeches, or that it’s important to drop bloodwormsinsidethe salamander
tank rather than next to it. Existing IL algorithms assume that the environment observed at test time
will be identical to the environment observed at training time, and so they cannot generalise to this
degree. Instead, we would like algorithms that solve the task ofrobust IL: given a small number of
demonstrations in one training environment, the algorithm should be able to generalise the intent
behind those demonstrations to (potentially very different) deployment environments.
One barrier to improved algorithms for robust IL is a lack of appropriate benchmarks. IL algorithms
are commonly tested on Reinforcement Learning (RL) benchmark tasks, such as those from OpenAI
Gym [ 37,23,27,8]. However, the demonstrator intent in these benchmarks is often trivial (e.g. the
goal for most of Gym’s MuJoCo tasks is simply to run forward), and limited variation in the initial
state distribution means that algorithms are effectively being evaluated in the same setting that was
used to provide demonstrations. Recent papers on Inverse Reinforcement Learning (IRL)—which
is a form of IL that infers a reward under which the given demonstrations are near-optimal—have
instead used “testing” variants of standard Gym tasks which differ from the original demonstration
environment [ 17,39,32,33]. For instance, Fu et al. [17]trained an algorithm on demonstrations from
the standard “Ant” task from Gym, then tested on a variant of the task where two of the creature’s
34th Conference on Neural Information Processing Systems (NeurIPS 2020), Vancouver, Canada.
Demonstrations…Colour test variant Layout test variantShape test variantFigure 1: Unlike existing IL benchmarks, MAGICAL makes a distinction betweendemonstration
andtestvariants of a task. Demonstrations are all provided in one particular conﬁguration of the
world (the “demonstration variant”). The learnt policy (or reward function) is then evaluated across a
set oftest variants, each of which randomise one aspect of the environment, such as block colour
or shape, environment layout, dynamics, etc. This makes it possible to understand precisely which
aspects of the underlying task the algorithm has been able to infer from demonstrations.
four legs were disabled. Splitting the environment into such “training” and “test” variants makes it
possible to measure the degree to which an algorithm overﬁts to task-irrelevant features of the supplied
demonstrations. However, there is so far no standard benchmark for robust IL, and researchers must
instead use ad-hoc adaptations of RL benchmarks—such as the modiﬁed Ant benchmark and similar
alternatives discussed in Section 5—to evaluate intent generalisation.
To address the above issues, we introduce the Multitask Assessment of Generalisation in Imitative
Control ALgorithms (MAGICAL). Each MAGICAL task occurs in the same 2D “MAGICAL uni-
verse”, where environments consist of a robot with a gripper surrounded by a variable number of
objects in aﬁxed-size workspace. Each task is associated with ademonstration variant, which is a
ﬁxed initial state from which all human demonstrations are provided. A task is also associated with a
set oftest variantsfor which no demonstrations are provided. As illustrated in Fig. 1, the test variants
each randomise a different aspect of the world, such as object colour, transition dynamics, or object
count. Randomising attributes of objects and the physics of the world lets us evaluate the ability of a
robust IL algorithm to performcombinatorial generalisation[ 5]. For instance, given a demonstration
of the robot pushing a red square across the workspace, an algorithm should be able to push a yellow
circle across the workspace; given a demonstration of three green and yellow blocks being placed in
a line, an algorithm should also be able to place four red and blue blocks in a line; and so on.
MAGICAL has several advantages over evaluation methods for standard (non-robust) IL:
•No “training on the test set”. Evaluating in the same setting that was used to give demon-
strations allows algorithms to exploit features that might not be present during deployment.
Having separate test variants for a task allows us to identify this kind of overﬁtting.
•Distinguishes between different types of transfer. Each test variant evaluates robustness
to a distinct, semantically meaningful axis of variation. This makes it possible to characterise
precisely which aspects of the provided demonstrations a given algorithm is relying on, and
to diagnose the causes of over- or under-ﬁtting.
•Enables knowledge reuse between tasks. Each MAGICAL task requires similar concepts
and low-level skills to solve. Different tasks can therefore provide “background knowledge”
for multi-task and meta-IL algorithms, such as knowledge that objects can have different
colours, or that objects with different shapes respond in a particular way when grasped.
Our experiments in Section 4 demonstrate the brittleness of standard IL algorithms, particularly
under large shifts in object position or colour. We also show that common methods for improv-
ing generalisation—such as multitask training, data augmentation, and alternative camera views—
sometimes improve robustness to small changes, but still fail to generalise to more extreme ones.
2 MAGICAL: Systematically evaluating robust IL
We will now introduce the main elements of the MAGICAL benchmark. Weﬁrst describe the abstract
setup of our benchmark, then detail the speciﬁc tasks and variants available in the benchmark.
2
2.1 Benchmark setup
The MAGICAL benchmark consists of a set oftasks T1,T2, . . . ,T m. Each task can in turn be broken
down intovariantsof a single base Markov Decision Process (MDP) that provide different state
distributions and “physics” for an agent. Formally, each task T= (S, vD,V) consists of ascoring
function S(τ) , ademonstration variant vD, and a set of additionaltest variants V={v 1, v2, . . . , v n}.
The scoring function S(τ) takes a trajectory τ= (s 0, a0, s1, a1, . . . , s T, aT)and assigns it a score
S(τ)∈[0,1] , where 0 is the score of a no-op policy, and 1 is the score of a perfect demonstrator.
Unlike a reward function, S(τ) need not be Markovian. In order to evaluate generalisation, the
variants are split into a single demonstration variantvDand a set of test variantsV.
In our domestic robotics analogy, vDmight represent a single room and time-of-day in which
demonstrations for some domestic task collected, while each test variant v∈V could represent
a different room, different time-of-day, and so on. Algorithms are expected to be able to take
demonstrations given only in demonstration variant vD, then generalise the intent behind those
demonstrations in order to achieve the same goal in each test variant v∈V . This can be viewed
either as a form of domain transfer, or as ordinary generalisation using only a single sample from a
hypothetical distribution over all possible variants of each task.
Formally, each variant associated with a task Tdeﬁnes a distribution over reward-free MDPs.
Speciﬁcally, a variant v= (p 0, pρ, H) consists of aninitial state distribution p0(s0), adynamics
distribution pρ(ρ), and ahorizon H. States are fully observable via an image-based observation
space. Further, all variants have the same state space, the same observation space, and the same
action space, which we discuss below. In addition to sampling an initial state s0∼p 0(s0), at the start
of each trajectory, adynamics vector ρ∈Rdis also sampled from the dynamics distribution pρ(ρ).
Unlike the state, ρis not observable to the agent; this vector controls aspects of the dynamics such as
friction and motor strength. Finally, the horizon Hdeﬁnes aﬁxed length for all trajectories sampled
from the MDP associated with the variantv. Each variant associated with a given task has the same
ﬁxed horizonHto avoid “leaking” information about the goal through early termination [27].
All tasks and variants in the MAGICAL benchmark share a common continuous state space S. A state
s∈S consists of a conﬁguration (pose, velocity, and gripper state) qRfor the robot, along with object
conﬁgurations O={o 1, . . . , o E}(where the number of objects in s0may be random). In addition to
pose, each object conﬁguration oiincludes an objecttypeand a number ofﬁxed attributes. Objects
can be of two types:blocks(small shapes that can be pushed around by the agent) andgoal regions
(coloured rectangles that the agent can drive over, but not push around). Each block has aﬁxed shape
(square, pentagon, star, or circle) and colour (red, green, blue, or yellow). Each goal region has aﬁxed
colour, width, and height. In order to facilitate generalisation across tasks with a different number of
objects, we use a common image-based observation space and discrete, low-level action space for
all tasks, which we describe in detail in Appendix A.1. At an implementation level, we expose each
variant of each task as a distinct Gym environment [ 8], which makes it straightforward to incorporate
MAGICAL into existing IL and RL codebases.
2.2 Tasks and variants
With the handful of building blocks listed in the previous section, we can create a wide variety of
tasks, which we describe in Section 2.2.1. The object-based structure of the environment also makes
it easy to evaluate combinatorial generalisation by randomising one or more attributes of each object
while keeping the othersﬁxed, as described in Section 2.2.2.
2.2.1 Tasks
Tasks in the MAGICAL suite were chosen to balance three desiderata. First, given a handful of
trajectories from the demonstration variant of a task, it should be possible for a human observer
to infer the goal with sufﬁcient accuracy to solve the test variants. We have chosen demonstration
variants (illustrated in Fig. 2) that rule out obvious misinterpretations, like mistakenly identifying
colour as being task-relevant when it is not. Second, the tasks should be constructed so that they
involve complementary skills that meta- and multi-task learning algorithms can take advantage of. In
our tasks, these “shared skills” include block manipulation; identiﬁcation of colour or shape; and
relational reasoning. Third, the demonstration variant of each task must be solvable by existing
(non-robust) IL algorithms. This ensures that the main challenge of the MAGICAL suite lies in
3
(a) MoveToCorner
 (b) MoveToRegion
 (c) MatchRegions
 (d) MakeLine
(e) FindDupe
 (f) FixColour
 (g) ClusterColour
 (h) ClusterShape
Figure 2: Demonstration variants for MAGICAL tasks. Appendix A shows an example demonstration
for each task.
generalisingto the test variants (robust IL), as opposed to reproducing the demonstrator’s behaviour
in the demonstration variant (standard IL). This section brieﬂy describes the resulting tasks; detailed
discussion of horizons, score functions, etc. is deferred to Appendix A.
Move to Corner (MTC) The robot must push a single block from one corner of the workspace to
the diagonally opposite corner. Test variants are constrained so that the robot and block start near the
lower right corner. The score is S(τ) = 1 if the blockﬁnishes the trajectory in the top left eighth of
the workspace, and decreases to zero as the block gets further from the top left corner.
MoveToRegion (MTR) The robot must drive inside a goal region and stay there. There are no
blocks in the demonstration or test variants. Further, variants only have one goal region to ensure that
the objective is unambiguous. The agent’s score is S(τ) = 1 if the robot’s body is inside the goal
region at the end of the trajectory, andS(τ) = 0otherwise.
MatchRegions (MR) There is a set of coloured blocks and a goal region visible to the robot, and
the robot must push all blocks of the same colour as the goal region into the goal region. Test variants
are constrained to have one goal region and at least one block of the same colour as that goal region.
A perfect score is given upon termination if the goal regions contains all and only blocks of the goal
region’s colour, with penalties for excluding any blocks of the goal colour, or including other blocks.
MakeLine (ML) Here the objective is for the robot to arrange all the blocks in the workspace into
a single line. A perfect score is given if all blocks are approximately colinear and close together;
a penalty is given for each block that does not form part of the longest identiﬁable line. Refer to
Appendix A for details on how a “line” is deﬁned.
FindDupe (FD) Similar to MatchRegions, except the goal region initially contains a “query” block
which has the same shape and colour as at least one other block outside the goal region. The objective
is to push at least one of those duplicate blocks into the goal region, which yields a perfect score.
Penalties are given for knocking the query block out of the goal region, failing toﬁnd a duplicate, or
pushing non-duplicate blocks into the goal region.
FixColour (FC) In each variant of this task, the workspace contains a set of non-overlapping goal
regions. Each goal region contains a single block, and exactly one block in the workspace will have a
different colour to its enclosing goal region. A perfect score is given for pushing that block out of its
enclosing goal region and into an unoccupied part of the workspace, without disturbing other blocks.
ClusterColour (CC) and ClusterShape (CS) The robot is confronted with a jumble of blocks of
different colours and shapes. It must push the blocks into clusters of either uniform colour (in the CC
4
task), or uniform shape (in the CS task). Test variants are constrained to include at least one block
of each colour and each shape. A perfect score is given for creating four spatially distinct clusters
corresponding to each of the four colours (CC) or shapes (CS), with a penalty proportional to the
number of blocks that do not belong to an identiﬁable cluster.
2.2.2 Test variants
In addition to its demonstration variant, each of the tasks above has a set of associated test variants.
Some variants are not supported for tasks that do not have any blocks, or where the initial state is
otherwise restricted, as documented in Table 2 of Appendix A.
Jitter Takes demo variant and randomly perturbs the poses of the robot and all objects by up to 5%
of the maximum possible range. Failure on this variant indicates severe overﬁtting to the
demonstration variant (e.g. by memorising action sequences).
Layout Completely randomises the position and orientation of the robot and all blocks, plus position
and dimensions of goal regions; a more challenging version of Jitter.
Colour Block colours are randomly reassigned as appropriate for the task. This tests whether the
agent is responsive to block colour (when it is task-relevant, like in CC and MR), or is
correctly ignorant of colour (when it is irrelevant, like in MTC and CS).
Shape Similar to Colour, except the shapes of blocks are randomised rather than the colours. This
variant either tests for appropriate responsiveness or invariance to shape, depending on
whether shape is task-relevant.
CountPlus The number of blocks is randomised (along with shape, colour, and position) to test
whether the agent can handle “larger” or “smaller” problems (i.e. “generalisation to n” [35]).
Dynamics Subtly randomises friction of objects and the robot against the workspace, as well as
force of robot motors (for rotation, forward/backward motion, and the gripper).
AllCombines all applicable variants for a task (e.g. Layout, Colour, Shape, CountPlus, Dynamics).
3 Data-efﬁcient intent disambiguation
Succeeding at the MAGICAL benchmark requires agents to generalise the intent behind a set of
demonstrations to substantially different test variants. We anticipate that resolving the ambiguity
inherent in this task will require additional sources of information about the demonstrator’s goal
beyond just single-task demonstrations. In this section, we review two popular non-robust IL
algorithms, as well as some common ways in which alternative sources of goal information are
incorporated into these algorithms to improve generalisation.
3.1 Baseline methods
Ourﬁrst baseline method is Behavioural Cloning (BC). BC treats a demonstration dataset Das
an undistinguished collection of state-action pairs {(s 1, a1), . . . ,(s M, aM)}. It then optimises the
parametersθof the policyπ θ(a|s)via gradient descent on the log loss
Lbc(θ;D) =−E
Dlogπ θ(a|s).
Our second baseline method is Generative Adversarial IL (GAIL) [ 23]. GAIL casts IL as a GAN
problem [ 19], where the generator πθ(a|s) is an imitation policy, and the discriminator Dψ:
S × A→[0,1] is tasked with distinguishing imitation behaviour from expert behaviour. Speciﬁcally,
GAIL uses alternating gradient descent to approximate a saddle point of
max
θmin
ψ�
Ladv(θ,ψ;D) =−E
πθlogD ψ(s, a)−E
Dlog(1−D ψ(s, a)) +λH(π θ)�
,
whereHdenotes entropy andλ≥0is a policy regularisation parameter.
We also included a slight variation on GAIL which (approximately) minimises Wasserstein divergence
between occupancy measures, rather than Jensen-Shannon divergence. We refer to this baseline as
5
WGAIL-GP. In analogy with WGAN-GP [20], WGAIL-GP optimises the cost
max
θmin
ψ�
Lw-gp(θ,ψ;D) =E
DDψ(s, a)−E
πθDψ(s, a) +λ w-gp E
1
2πθ+1
2D(�∇sD(s, a)� 2−1)2�
,
The gradient penalty approximately enforces 1-Lipschitzness of the discriminator by encouraging the
norm of the gradient to be 1 at points between the support of πθandD. Since actions were discrete,
we did not enforce 1-Lipschitzness with respect to the action input. We also did not backpropagate
gradients with respect to the gradient penalty back into the policy parameters θ, since the gradient
penalty is only intended as a soft constraint onD.
In addition to these baselines, we also experimented with Apprenticeship Learning (AL). Unfortu-
nately we could not get AL to perform well on most of our tasks, so we defer further discussion of
AL to Appendix B.
3.2 Using multi-task data
As noted earlier, the MAGICAL benchmark tasks have similar structure, and should in principle
beneﬁt from multi-task learning. Speciﬁcally, say we are given a multi-task dataset Dmt=
{D(T i, vD
i, ni)}M
i=1, where D(T i, v, n) denotes a dataset of ntrajectories for variant vof task Ti.
For BC and GAIL, we can decompose the policy for task Tiasπi
θ=gi
θ◦fθ, where fθ:S→Rd
is a multi-task state encoder, while gi
θ:Rd→Δ(A) is a task-speciﬁc policy decoder. We can
also decompose the GAIL discriminator as Di
ψ=si
ψ◦rψ, where rψ:S×A→Rdis shared and
si
ψ:Rd→[0,1]is task-speciﬁc. We then modify the BC and GAIL objectives to
Lbc(θ;D mt) =M�
i=1Lbc(θ;D(T i, vD
i, ni))andL adv(θ,ψ;D mt) =M�
i=1Ladv(θ,ψ;D(T i, vD
i, ni)).
3.3 Domain-speciﬁc priors and biases
Often the most straightforward way to improve the robustness of an IL algorithm is to constrain the
solution space to exclude common failure modes. For instance, one could use a featurisation that
only captures task-relevant aspects of the state. Such priors and biases are generally domain-speciﬁc;
for the image-based MAGICAL suite, we investigated two such biases:
•Data augmentation: In MAGICAL, our score functions are invariant to whether objects
are repositioned or rotated slightly; further, human observers are typically invariant to small
changes in colour or local image detail. As such, we used random rotation and translation,
Gaussian noise, and colour jitter to augment training data for the BC policy and GAIL
discriminator. This can be viewed as a post-hoc form of domain randomisation, which
has previously yielded impressive results in robotics and RL [ 2]. We found that GAIL
discriminator augmentations were necessary for the algorithm to solve more-challenging
tasks, as previously observed by Zolna et al. [41]. In BC, we found that policy augmentations
improved performance on both demonstration and test variants.
•Ego- and allocentric views: Except where indicated otherwise, all of the experiments
in Section 4 use anegocentricperspective, which always places the agent at the same
position (and in the same orientation) within the agent’sﬁeld of view. This contrasts with an
allocentricperspective, where observations are focused on aﬁxed region of the environment
(in our case, the extent of the workspace), rather than following the agent’s position. In
the context of language-guided visual navigation, Hill et al. [22]previously found that an
egocentric view improved generalisation to unseen instructions or unseen visual objects,
despite the fact that it introduces a degree of partial observability to the environment.
4 Experiments
Our empirical evaluation has two aims. First, to conﬁrm that single-task IL methods fail to generalise
beyond the demonstration variant in the MAGICAL suite. Second, to analyse the ways in which the
common modiﬁcations discussed in Section 3 affect generalisation.
6
Method Demo Jitter Layout Colour Shape
BC (single-task) 0.64±0.29 0.56±0.27 0.14±0.16 0.39±0.30 0.52±0.33
Allocentric 0.58±0.33 0.48±0.29 0.04±0.04 0.42±0.32 0.50±0.37
No augmentations 0.55±0.37 0.37±0.30 0.12±0.15 0.33±0.30 0.41±0.33
No trans./rot. aug. 0.55±0.37 0.41±0.31 0.13±0.15 0.33±0.30 0.43±0.35
Multi-task 0.59±0.33 0.53±0.31 0.14±0.18 0.30±0.25 0.51±0.36
GAIL (single-task) 0.72±0.35 0.69±0.33 0.22±0.23 0.27±0.24 0.60±0.42
Allocentric 0.57±0.46 0.49±0.40 0.03±0.03 0.39±0.36 0.50±0.45
No augmentations 0.44±0.42 0.32±0.31 0.09±0.12 0.19±0.23 0.28±0.33
WGAIL-GP 0.42±0.38 0.33±0.32 0.14±0.20 0.10±0.11 0.33±0.33
Multi-task 0.37±0.41 0.33±0.36 0.16±0.25 0.11±0.12 0.28±0.36
Table 1: Score statistics for a subset of variants and compared algorithms. We report the mean and
standard deviation of test scores aggregated acrossalltasks, withﬁve seeds per algorithm and task.
Darker colours indicate higher scores.
4.1 Experiment details
We evaluated all the single- and multi-task algorithms in Section 3, plus augmentation and perspective
ablations, on all tasks and variants. Each algorithm was trainedﬁve times on each task with different
random seeds. In each run, the training dataset for each task consisted of 10 trajectories from the
demo variant. All policies, value functions, and discriminators were represented by Convolutional
Neural Networks (CNNs). Observations were preprocessed by stacking four temporally adjacent RGB
frames and resizing them to 96 ×96 pixels. For multi-task experiments, task-speciﬁc weights were
used for theﬁnal fully-connected layer of each policy/value/discriminator network, but weights of all
preceding layers were shared. The BC policy and GAIL discriminator both used translation, rotation,
colour jitter, and Gaussian noise augmentations by default. The GAIL policy and value function
did not use augmented data, which we found made training unstable. Complete hyperparameters
and data collection details are listed in Appendix B. The IL algorithm implementations that we used
to generate these results are available on GitHub,1as is the MAGICAL benchmark suite and all
demonstration data.2
4.2 Discussion
Due to space limitations, this section addresses only a selection of salient patterns in the results.
Table 1 provides score statistics for a subset of algorithms and variants, averaged acrossalltasks.
See Section 2.2.1 for task name abbreviations (MTR, FC, etc.). Because the tasks vary in difﬁculty,
pooling across all tasks yields high score variance in Table 1. Actual score variance for each method
is much lower when results are constrained to just one task; refer to Appendix C for complete results.
Overﬁtting to position All algorithms exhibited severe overﬁtting to thepositionof objects. The
Layout, CountPlus, and All variants yielded near-zero scores in all tasks except MTC and MTR,
and on many tasks there was also poor transfer to the Jitter variant. For some tasks, we found that
the agent would simply execute the same motion regardless of its initial location or the positions of
task-relevant objects. This was true on the FC task, where the agent would always execute a similar
forward arc regardless of its initial position, and also noticeable on MTC and FD, where the agent
would sometimes move to the side of a desired block when it was shifted slightly. For BC, this issue
was ameliorated by the use of translation and rotation augmentations, presumably because the policy
could better handle small deviations from the motions seen at training time.
Colour and shape transfer Surprisingly, BC and GAIL both struggled with colour transfer to a
greater degree than shape transfer on several tasks, as evidenced by the aggregated statistics for
Colour and Shape variants in Table 1. Common failure modes included freezing in place or moving
in the wrong direction when confronted with an object of a different colour to that seen at training
time. In contrast, in most tasks where shape invariance was desirable (including MTC, MR, ML, and
FC), the agent had no trouble reaching and manipulating blocks of different shapes. Although colour
1Multi-task imitation learning algorithms:https://github.com/qxcv/mtil/
2Benchmark suite and links to data:https://github.com/qxcv/magical/
7
jitter was one of the default augmentations, the BC ablations in Table 1 suggest that almost all of the
advantage of augmentations comes from the use of translation/rotation augmentations. In particular,
we did notﬁnd that colour jitter greatly improved performance on tasks where the optimal policy
was colour-invariant. In spite of exposing the networks to a greater range of colours at train time,
multitask training also failed to improve colour transfer, as we discuss below. Although translation
and rotation sometimes improved colour transfer (e.g. for BC on FindDupe in Table 7), it is not clear
why this was the case. We speculate that these augmentations could have encouraged the policy to
acquire more robust early-layer features for edge and corner detection that did not rely on just one
colour channel.
Multi-task transfer Plain multi-task learning had mixed effects on generalisation. In some cases
it improved generalisation (e.g. for BC on FC), but in most cases it led to unchanged ornegative
transfer, as in the Colour test variants for MTC, MR, and FD. This could have been because the policy
was using colour to distinguish between tasks. More speculatively, it may be that a multi-task BC
or GAIL loss is not the best way to incorporate off-task data, and that different kinds of multi-task
pretraining are necessary (e.g. learning forward or inverse dynamics [9]).
Egocentric view and generalisation The use of an allocentric (rather than egocentric) view did
not improve generalisation or demo variant performance for most tasks, and sometimes decreased
it. Table 1 shows the greatest performance drop on variants that change object position, such as
Layout and Jitter. For example, in MTR we found that egocentric policies tended to rotate in one
direction until the goal region was in the centre of the agent’sﬁeld of view, then moved forward
to reach the region, which generalises well to different goal region positions. In contrast, the
allocentric policy would often spin in place or get stuck in a corner when confronted with a goal
region in a different position. This supports the hypothesis of Hill et al. [22]that the egocentric view
improves generalisation by creating positional invariances, and reinforces the value of being able to
independently measure generalisation across distinct axes of variation (position, shape, colour, etc.).
5 Related work
There are few existing benchmarks that speciﬁcally examine robust IL. The most similar benchmarks
to MAGICAL have appeared alongside evaluations of IRL and meta-IL algorithms. As noted in
Section 1, several past papers employ “test” variants of standard Gym MuJoCo environments to
evaluate IRL generalisation [ 17,39,32,33], but these modiﬁed environments tend to have trivial
reward functions (e.g. “run forward”) and do not easily permit cross-environments transfer. Xu
et al. [38]and Gleave and Habryka [18]use gridworld benchmarks to evaluate meta- and multi-task
IRL, and both benchmarks draw a distinction between demonstration and execution environments
within a meta-testing task. This distinction is similar in spirit to the demonstration/test variant split in
MAGICAL, although MAGICAL differs in that it has more complex tasks and the ability to evaluate
generalisation across different axes. We note that there also exist dedicated IL benchmarks [30, 26],
but they are aimed at solving challenging robotics tasks rather than evaluating generalisation directly.
There are many machine learning benchmarks that evaluate generalisation outside of IL. For instance,
there are several challenging benchmarks for generalisation [ 31,12,13] and meta- or multi-task
learning [ 40] in RL. Unlike MAGICAL, these RL benchmarks have no ambiguity about what the goal
is in the training environment, since it is directly speciﬁed via a reward function. Rather, the challenge
is to ensure that the learnt policy (for model-free methods) can achieve that clearly-speciﬁed goal
in different contexts (RL generalisation), or solve multiple tasks simultaneously (multi-task RL), or
be adapted to new tasks with few rollouts (meta-RL). There are also several instruction-following
benchmarks for evaluating generalisation in natural language understanding [ 29,34]. Although
these are not IL benchmarks, they are similar to MAGICAL in that they include train/test splits that
systematically evaluate different aspects of generalisation. Finally, the Abstract Reasoning Corpus
(ARC) is a benchmark that evaluates the ability of supervised learning algorithms to extrapolate
geometric patterns in a human-like way [ 11]. Although there is no sequential decision-making
aspect to ARC, Chollet [11]claims that solving the corpus may still require priors for “objectness”,
goal-directedness, and various geometric concepts, which means that methods suitable for solving
MAGICAL may also be useful on ARC, and vice versa.
Although we covered some simple methods of improving IL robustness in Section 3, there also exist
more sophisticated methods tailored to different IL settings. Meta-IL [ 15,25] and meta-IRL [ 38,39]
8
algorithms assume that a large body of demonstrations is available for some set of “train tasks”,
but only a few demonstrations are available for “test tasks” that might be encountered in the future.
Each test task is assumed to have a distinct objective, but one that shares similarities with the train
tasks, making it possible to transfer knowledge between the two. These methods are likely useful for
multi-task learning in the context of MAGICAL, too. However, it’s worth noting that past meta-IL
work generally assumes that meta-train and meta-test settings are similar, whereas this work is
concerned with how to generalise the intent behind a few demonstrations given in one setting (the
demo variant) to other, potentially very different settings (the test variants). Similar comments apply
to existing work on multi-task IL and IRL [18, 10, 14, 3].
6 Conclusion
In this paper, we introduced the MAGICAL benchmark suite, which is theﬁrst imitation learning
benchmark capable of evaluating generalisation across distinct, semantically-meaningful axes of
variation in the environment. Unsurprisingly, results for the MAGICAL suite conﬁrm that single-task
methods fail to transfer to changes in the colour, shape, position and number of objects. However, we
also showed that image augmentations and perspective shifts only slightly ameliorate this problem, and
multi-task training can sometimes make itworse. This lack of generalisation stands in marked contrast
to human imitation: even 14-month-old infants have been observed to generalise demonstrations
of object manipulation tasks across changes in object colour and shape, or in the appearance of the
surrounding room [ 4]. Closing the gap between current IL capabilities and human-like few-shot
imitation could require signiﬁcant innovations in multi-task learning, action and state representations,
or models of human cognition. The MAGICAL suite provides a way of evaluating such algorithms
which not only tests whether they generalise well “on average”, but also shines a light on the speciﬁc
kinds of generalisation which they enable.
7 Broader impact
This paper presents a new benchmark for robust IL and argues for an increased focus on algorithms
that can generalise demonstrator intent across different settings. We foresee several possible follow-on
effects from improved IL robustness:
Economic effects of automation Better IL generalisation could allow for increased automation in
some sectors of the economy. This has the positiveﬂow-on effect of increased economic
productivity, but could lead to socially disruptive job loss. Because our benchmark focuses
on robust IL in robotics-like environments, it’s likely that any effect on employment would
be concentrated in sectors involving activities that are expensive to record. This could
include tasks like surgery (where few demonstrators are qualiﬁed to perform the task, and
privacy considerations make it difﬁcult to collect data) or packaging retail goods for postage
(where few-shot learning might be important when there are many different types of goods
to handle).
Identity theft and model extraction More robust IL could enable better imitation ofspeciﬁcpeople,
and not just imitation of people in general. This could lead to identify theft, for instance by
mimicking somebody’s speech or writing, or by fooling biometric systems. Because this
benchmark focuses on control and manipulation rather than media synthesis, it’s unlikely
that algorithms designed to solve our benchmark will be immediately useful for this purpose.
On the other hand, this concern is still relevant when applied to machine behaviour, rather
than human behaviour. In NLP, it’s known that weights for ML models can be “stolen” by
observing the model’s outputs for certain carefully chosen inputs [ 28]. Similarly, more
robust IL could make it possible to clone a robot’s policy by observing its behaviour, which
could make it harder to sell robot control algorithms as standalone products.
Learnt objectives Hadﬁeld-Menell et al. [21] argues that it is desirable for AI systems to infer
their objectives from human behaviour, rather than taking them asﬁxed. This can avoid
problems that arise when an agent (human, robot, or organisation) doggedly pursues an easy-
to-measure but incorrect objective, such as a corporate executive optimising for quarterly
proﬁt (which is easy to measure) over long-term proﬁtability (which is actually desired by
shareholders). IL makes it possible to learn objectives from observed human behaviour, and
9
more robust IL may therefore lead to AI systems that better serve their designers’ goals.
However, it’s worth noting that unlike, say, HAMDPs [ 16] or CIRL games [ 21], IL cannot
request clariﬁcation from a demonstrator if the supplied demonstrations are ambiguous,
which limits its ability to learn the right objective in general. Nevertheless, we hope that
insights from improved IL algorithms will still be applicable to such interactive systems.
Acknowledgments and Disclosure of Funding
We would like to thank reviewers for helping to improve the presentation of the paper (in particular,
clarifying the distinction between traditional IL and robust IL), and for suggesting additional related
work and baselines. This work was supported by a Berkeley Fellowship and a grant from the Open
Philanthropy Project.
References
[1]Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning.
InICML, 2004.
[2]Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur
Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, et al. Solving Rubik’s
Cube with a robot hand.arXiv preprint arXiv:1910.07113, 2019.
[3]Monica Babes, Vukosi Marivate, Kaushik Subramanian, and Michael L Littman. Apprenticeship
learning about multiple intentions. InICML, 2011.
[4]Sandra B Barnat, Pamela J Klein, and Andrew N Meltzoff. Deferred imitation across changes
in context and object: Memory and generalization in 14-month-old infants.Infant Behavior &
Development, 19(2):241, 1996.
[5]Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius
Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan
Faulkner, et al. Relational inductive biases, deep learning, and graph networks.arXiv preprint
arXiv:1806.01261, 2018.
[6]Aude Billard, Sylvain Calinon, Ruediger Dillmann, and Stefan Schaal. Survey: Robot program-
ming by demonstration. In Bruno Siciliano and Oussama Khatib, editors,Handbook of robotics,
chapter 59. Springer, 2008.
[7]Robert C Bolles and Martin A Fischler. A RANSAC-based approach to modelﬁtting and its
application toﬁnding cylinders in range data. InIJCAI, volume 1981, pages 637–643, 1981.
[8]Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang,
and Wojciech Zaremba. OpenAI Gym.arXiv:1606.01540, 2016.
[9]Daniel S Brown, Russell Coleman, Ravi Srinivasan, and Scott Niekum. Safe imitation learning
via fast bayesian reward inference from preferences.arXiv:2002.09089, 2020.
[10] Jaedeug Choi and Kee-Eung Kim. Nonparametric bayesian inverse reinforcement learning for
multiple reward functions. InNIPS, pages 305–313, 2012.
[11] François Chollet. The measure of intelligence.arXiv preprint arXiv:1911.01547, 2019.
[12] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying
generalization in reinforcement learning.arXiv:1812.02341, 2018.
[13] Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural
generation to benchmark reinforcement learning.arXiv:1912.01588, 2019.
[14] Christos Dimitrakakis and Constantin A Rothkopf. Bayesian multitask inverse reinforcement
learning. InEuropean Workshop on Reinforcement Learning. Springer, 2011.
10
[15] Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAI Jonathan Ho, Jonas Schneider, Ilya
Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. InNeurIPS,
pages 1087–1098, 2017.
[16] Alan Fern, Sriraam Natarajan, Kshitij Judah, and Prasad Tadepalli. A decision-theoretic model
of assistance. InIJCAI, 2007.
[17] Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse
reinforcement learning.arXiv preprint arXiv:1710.11248, 2017.
[18] Adam Gleave and Oliver Habryka. Multi-task maximum entropy inverse reinforcement learning.
arXiv preprint arXiv:1805.08882, 2018.
[19] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil
Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. InNIPS, 2014.
[20] Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville.
Improved training of Wasserstein GANs. InNIPS, 2017.
[21] Dylan Hadﬁeld-Menell, Stuart J Russell, Pieter Abbeel, and Anca Dragan. Cooperative inverse
reinforcement learning. InNIPS, pages 3909–3917, 2016.
[22] Felix Hill, Andrew K Lampinen, Rosalia Schneider, Stephen Clark, Matthew Botvinick, James L
McClelland, and Adam Santoro. Environmental drivers of systematicity and generalization in a
situated agent. InICLR, 2020.
[23] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. InNIPS, 2016.
[24] Jonathan Ho, Jayesh Gupta, and Stefano Ermon. Model-free imitation learning with policy
optimization. InICML, 2016.
[25] Stephen James, Michael Bloesch, and Andrew J Davison. Task-embedded control networks for
few-shot imitation learning.CORL, 2018.
[26] Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew J Davison. Rlbench: The robot
learning benchmark & learning environment.arXiv preprint arXiv:1909.12271, 2019.
[27] Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan
Tompson. Discriminator-actor-critic: Addressing sample inefﬁciency and reward bias in
adversarial imitation learning. InICLR, 2019.
[28] Kalpesh Krishna, Gaurav Singh Tomar, Ankur P Parikh, Nicolas Papernot, and Mohit Iyyer.
Thieves on Sesame Street! Model extraction of BERT-based APIs.ICLR, 2020.
[29] Brenden M Lake and Marco Baroni. Generalization without systematicity: On the compositional
skills of sequence-to-sequence recurrent networks.ICML, 2018.
[30] Raphael Memmesheimer, Ivanna Mykhalchyshyna, Viktor Seib, and Dietrich Paulus. Simitate:
A hybrid imitation learning benchmark.arXiv preprint arXiv:1905.06002, 2019.
[31] Alex Nichol, Vicki Pfau, Christopher Hesse, Oleg Klimov, and John Schulman. Gotta learn fast:
A new benchmark for generalization in RL.arXiv:1804.03720, 2018.
[32] Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine. Variational
discriminator bottleneck: Improving imitation learning, inverse RL, and GANs by constraining
informationﬂow.arXiv preprint arXiv:1810.00821, 2018.
[33] Ahmed H Qureshi, Byron Boots, and Michael C Yip. Adversarial imitation via variational
inverse reinforcement learning.ICLR, 2019.
[34] Laura Ruis, Jacob Andreas, Marco Baroni, Diane Bouchacourt, and Brenden M Lake. A bench-
mark for systematic generalization in grounded language understanding.arXiv:2003.05161,
2020.
11
[35] Iude W Shavlik. Acquiring recursive and iterative concepts with explanation-based learning.
Machine Learning, 1990.
[36] Adam Stooke and Pieter Abbeel. rlpyt: A research code base for deep reinforcement learning in
pytorch.arXiv preprint arXiv:1909.01500, 2019.
[37] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. InIJCAI,
2018.
[38] Kelvin Xu, Ellis Ratner, Anca Dragan, Sergey Levine, and Chelsea Finn. Learning a prior over
intent via meta-inverse reinforcement learning.ICML, 2019.
[39] Lantao Yu, Tianhe Yu, Chelsea Finn, and Stefano Ermon. Meta-inverse reinforcement learning
with probabilistic context variables. InNeurIPS, pages 11749–11760, 2019.
[40] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and
Sergey Levine. Meta-World: A benchmark and evaluation for multi-task and meta reinforcement
learning. InCoRL, 2020.
[41] Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarej, David Budden,
Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation
learning.arXiv preprint arXiv:1910.01077, 2019.
12