Abstract: Iterated Distillation and Amplification (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 highschool mathematics and investigates whether learning to decompose problems can improve performance over supervised learning. 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 computation time (varying compute at inference time as a function of the input) can improve the robustness and efficiency 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.

Owain Evans
William Saunders
Andreas Stuhlmüller

Machine Learning Projects for Iterated Distillation and Amplification

Iterated Distillation and Amplification (IDA) is a framework for training models from data  [1, 2] . IDA is related to and builds on existing frameworks like supervised 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 superhuman performance at the task.  1  3. Humans have some high-level understanding of the task and can also provide 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 fixed 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] .
 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 T contains a series of tasks that get progressively harder for humans to solve. Formally, let T = N i=0 T i , where tasks in T n are harder than tasks in T n−1 for all n. We are given a training set which includes solutions to the easiest class of problems T 0 and human demonstrations of decomposing tasks T n into finitely many slightly easier tasks in T n−1 . In IDA, the initial training steps are: 1. M is trained by supervised learning 2 to reproduce the answers to the easiest tasks T 0 . 2. M is trained by supervised learning to imitate the human demonstrations for decomposing a task x ∈ T n into a set of tasks in T n−1 and 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 the recursive step. After initial training, M can solve tasks in T 1 by first decomposing them into tasks in T 0 . M is then trained by supervised learning on its own solutions to tasks in T 1 , enabling M to solve T 1 tasks directly (i.e. without decomposition). Solving a task in T 1 directly involves a single call to M , while solving by decomposition into tasks in T 0 requires M to be called on each of the T 0 tasks. This process of "training on its own solutions" can be iterated. M is 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 M can already solve. If supervised training works perfectly at each iteration, then eventually M can solve any task in T directly (with only single call to M ). (This depends on the strong assumption that humans can decompose all tasks T n into tasks in T n−1 -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):
 • Amplification Step M solves tasks in T n by decomposing them into tasks in T n−1 , which it solves directly (without decomposition).
 • Distillation Step M is trained by supervised learning to solve tasks in T n directly, with target solutions coming from the Amplification Step.
 It is called the "Amplification Step" because it amplifies the capability of model M . While M can only solve tasks in T n−1 directly, M can be used (via decomposition and aggregation) to solve tasks in T n . In the "Distillation Step", the slower amplified model (which makes multiple calls to M ) is distilled 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 M directly and when to fall back to the amplified model (see Project 3). Second, M solves each task x ∈ T1 by decomposing into tasks {xa, x b } ∈ T0 and solving these directly. This eventually produces a dataset of solved tasks x ∈ T1, which M is trained to solve directly by supervised learning (leftmost "distill" step). This process is then repeated. x ∈ T 0 x a ∈ T 0 amplify distill amplify x ∈ T 1 x b ∈ T 0 x ∈ T 1 x a ∈ T 1 x ∈ T 2 x b ∈ T 1 distill x ∈ T 2 amplify x a ∈ T 2 x ∈ T 3 x b ∈ T 2 To simplify the exposition, we have presented M as being trained to decompose tasks and learn all base case solutions before the iterative training by Amplification 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] .
 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.  3  These examples are intended to build intuition for IDA. They are not themselves practically relevant use cases for IDA.
 Example 1: Multiplication This toy example shows how one would use IDA to train a neural net M to multiply large integers. We assume that M has 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). M is trained as follows: 1. Train M on simple multiplications: single-digit multiplication and multiplication by 10. For example: 5 × 6 = 30, 10 × 234 = 2340, 9 × 8 = 72.  Recursing all the way to the base case would require many calls to the neural net M . Instead, M could be trained incrementally using the IDA scheme described above. If this was successful, M could eventually compute 384 × 19 or 79332 × 2927 in 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 M that 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) , m × n) for m, n < 10 6 . Training by IDA will generally make learning more difficult, because M will be trained by supervised learning on its own answers to multiplication problems (instead of ground-truth answers).
 Train 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 M into 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 algorithm to distill an exponential game-tree expansion into a feedforward neural 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 find, wildcard search, and shortest path. To apply IDA to finding the shortest path between nodes s and t in 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 s and t are adjacent). 2. A set of demonstrations of decomposing shortest path problems into smaller shortest path problems and aggregating the results. The decomposition in (  2 ) is similar 4 to the following recursion, which is also used in the dynamic programming algorithm for shortest path: min-dist(s, t) = min ({min-dist(s, x) + dist(x, t) | x adjacent to t}) Here "min-dist(s, t)" is the minimum path length between nodes s and t, and "dist" is the distance between adjacent nodes. Christiano et al. use a Transformer model  [13]  as their model M and they compare two ways of training the Transformer to compute the shortest path. The first 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.
 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 using 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 approximation, 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 systems, it would likely need to apply to tasks (and use decompositions) that involve sophisticated reasoning in natural language. There is currently no published 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 different incentive structures for IDA-like approaches (e.g. different objective functions and reward signals) lead to aligned behavior. See examples from Ought  [20]  and OpenAI  [21]  and stay in touch with Ought 5 for the latest information on experiments. 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 compatible 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 computational complexity theory, and describes links between Debate and IDA.
 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 different ways: • The projects train a single model M to perform a wide range of types of decompositions. By contrast, Christiano et al. train a fresh model for each algorithmic problem (i.e. one for each of union find, shortest path, 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 generalization 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 amplification. Improving test-time performance by amplification 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 reinforcement learning systems. • Connect IDA to existing work in ML. The projects below connect to language modeling, neural programming, calibration for neural networks, and adaptive computation time. IDA is also related to semi-supervised learning, deep reinforcement learning, dynamic programming, belief propagation, etc. A research project could explore and develop any of these connections. • 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 community. The projects below are not primarily aimed to achieve this, but they may provide a useful first 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 domains (e.g. common-sense reasoning) where we are not able to write down an algorithm to decompose problems. 1 Project 1: Amplifying Mathematical Reasoning
 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 first sub-problem decomposes to substituting the function g and expanding the resulting expression.
 Project Directions The aim of this project is to train a model via IDA to solve mathematics problems. Following Christiano et al  [8] , we could use IDA at training time to 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 amplification at test time to apply more compute 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. Advanced mathematics problems: competition or Olympiad math, universitylevel 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 gametree 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.
 Write an algorithm that solves problems by decomposition as in Christiano 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 mathematics. However, it is unlikely that the usual way people solve problems would yield the most useful decompositions. So, part of the research effort 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 first kind is an efficient 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. Experiments 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 inefficiently 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 efficient neural algorithm. However, for advanced mathematics problems in natural language, it is not clear what this inefficient algorithm would look like.  6  So there is a separate research project in investigating such algorithms. In the next section, we outline a project that uses high-school math problems in natural language. We suggest starting with algorithmically generated decompositions and later extending this to decompositions provided by humans. The problems are generated by an algorithm, so there is unlimited training data. The dataset includes test sets for both interpolation and extrapolation. The extrapolation questions have quantities that vary outside of the range encountered on the training set.
 IDA for High School Mathematics 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 Transformer with 30M parameters, which achieves 76% model accuracy (probability of a correct answer) on interpolation and 50% on extrapolation.
 Approach with IDA How should one generate decompositions for training IDA to solve these math problems? One option is to write an algorithm for decomposing problems, rather than collecting human decompositions. The problems were generated using a compositional algorithm: consider the last example in the list above, which 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 fit for IDA training. (For instance, the best decompositions could be more or less fine-grained.) One objective would be to do better than the supervised baseline at test time by applying amplification (i.e. decomposing the problem and using multiple calls to the model M ). In particular, amplification 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.
 Non-Amplification 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 map strings to strings. Another baseline would make use of the same decomposition 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 approach 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).
 Questions to investigate The project would seek to investigate some of the following questions: • What is the test-time performance of applying amplification 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 off-distribution problems? • 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 different kinds of decomposition strategies? • How robust is the amplified model to noise in the training data and to approximation error in the neural net? • Can we find 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 distribution over sub-problem examples. How does this impact performance? Are there principles for generating training data in the most useful way for IDA? • Using amplification 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 algorithmically generated decompositions? 1.4 Related Work • 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 solving 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 for Machine Learning of Higher-Order Theorem Proving (Bansal et al.)  [23] . This dataset contains fully formalized proofs for a large number of theorems and a framework for training ML systems to produce proofs. 2 Project 2: IDA for Neural Program Interpretation
 Motivation Many algorithms (e.g. matrix arithmetic, shortest path, sorting) involve decomposing 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 operations of the language. We can explore the capabilities of IDA by training a model on the decompositions used in the execution of computer programs. Training IDA on decompositions from programs is closely related to research on Neural Program Interpretation (see Reed and de Freitas  [28]  and related work 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 decomposition (i.e. the way procedures call other procedures). As with IDA, one motivation for learning this internal behavior is to achieve stronger generalization. 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] ).
 Project Directions The project could focus on either of the following areas:
 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 efficient "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 decompositions Much of the NPI work uses environments and architectures designed specially 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 approximation error-and errors will usually be larger for inputs that are off the training distribution. This will cause problems both during IDA's iterative training procedure and also at test time. Part of the project would be to investigate how well distillation works for different kinds of programs and for different 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 using amplification. Project 3 (below) explores this "adaptive computation" or "meta-reasoning" approach.
 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 operate on human-readable structures like natural languages or images. We also think that pure functional programming is a better fit 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, we might not know about these "basic operations on data" 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 most of 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.
 Non-Amplification 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 significant overhead and so is unlikely to *improve* performance. One could also consider tasks that require ML to process the inputs (e.g. tasks involving images). The baseline in this case would be a program that defers some decisions to a classifier.
 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 approach 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 supervision 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 inference core for executing three different programs (addition, sorting, rotating CAD models). So one set of LSTM parameters are used to execute all programs. However, the different programs are stored as different embeddings, 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 specific 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 off 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. Adaptive Neural Compilation (Bunel et al.)  [34]  Takes programs, translates them into a differentiable form, then uses backpropagation to optimize them. They optimize programs to be more efficient on a restricted training distribution of problems. Recent Advances in Neural Program Synthesis (Kant)  [35]  This paper provides a summary of different approaches to NPI and Neural Program Synthesis. Learning Compositional Neural Programs with Recursive Tree Search and Planning (Pierrot et al.)  [36]  NPI approach that "incorporates the strengths of Neural Programmer-Interpreters (NPI) and AlphaZero". While relevant to IDA, this approach differs in important ways. 3 Project 3: Adaptive Computation
 Motivation An ML model exhibits "adaptive computation" if it intelligently varies its computations for different 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 fixed 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 "metareasoning". 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 learning  [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 current deep learning research. Yet human cognition is often adaptive: people dynamically decide how much time to spend on a task based on its perceived difficulty. 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 profile: 1. Training time: Large amounts of compute and time are permitted. For example, training might take months and use many CPUs and GPUs.
 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 profile. When AlphaZero  [3]  plays competitive matches, it has time to make many calls to a neural net for each move. Alp-haZero uses MCTS to investigate promising moves from the current position. As ML is applied to more tasks for which humans spend a long time thinking (e.g. mathematics, science, business strategy), the profile 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, produced 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 decomposition strategy ("Amplification") and involves making multiple calls to M .
 Project Directions Adaptive computation for IDA can be applied either during the iterative training scheme or at test/deployment time:
 Training time The goal during training is to bootstrap by using amplification to solve progressively harder problems. Adaptive computation time could be applied 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 amplified process (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.
 Test time The goal is to perform well on a set of test problems given a time budget. 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 benefit from more compute time (via amplification). There are also more fine-grained questions about how to use the compute budget. For example, when using amplification, which problems should be decomposed first and how far should the recursive decomposition go?
 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 decides between a fast (distilled) process for solving problems and a slow process that uses a fixed amount of amplification. The fast process makes a single call to a neural net, while the slow process makes at most n calls to the same net. For each input problem, the algorithm runs the fast process and has to decide whether to also run the slow process. The goal is to balance the greater accuracy of the slow process with a time/compute cost that is specified as part of the problem statement. This decision of whether to run the slow process could be trained end-toend, by always running both processes during training (as in  [37] ). A different approach would be to use a calibrated model for the fast process and decide whether to call the slow process based on the confidence of this calibrated model. 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 sufficient 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 differ 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 fixed 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.
 Non-Amplification Baselines For any adaptive computation approach, it is important to compare against nonadaptive approaches (which use a fixed amount of compute) and also against simple heuristics for scaling up compute for harder instances. Amplification could also be compared to adaptive approaches that do not use amplification (e.g. something like  [37] ).
 Related Work Attending to Mathematical Language with Transformers (Wangperawong)  [44]  There is a dataset (available at https://github.com/tensorflow/tensor2tensor) and paper for addition/multiplication/subtraction of numbers. The paper explores Transformers that can use more compute for larger instances. Adaptive Computation Time for Recurrent Neural Networks (Graves)  [37]  Adaptive Computation Time for RNNs adds the probability of halting computation at current step to the output of RNN. The idea is to train the adaptive RNN end-to-end: running long computations at training time, which allows differentiation 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 fixed number of steps between predictions). Universal Transformers (Dehghani et al.)  [46]  The paper applies Adaptive Computation Time to the Transformer architecture. It shows improved performance on bAbI tasks with Adaptive Computation Time On Calibration of Modern Neural Networks (Guo et al.)  [47]  This paper tests the calibration of convolutional neural networks and finds 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 tradeoff 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 contributions 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 endto-end, by playing single-player gridworld game Sokoban. They learn a more sample efficient search than regular MCTS. Contents0a 14 3 18 0 1418 Background on IDA 0.1 What is IDA? . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 0.2 Examples . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 0.3 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 0.4 Project Goals . . . . . . . . . . . . . . . . . . . . . . . . . . . . .1 Project 1: Amplifying Mathematical Reasoning 1.1 Motivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1.2 Project Directions . . . . . . . . . . . . . . . . . . . . . . . . . . 1.3 IDA for High School Mathematics . . . . . . . . . . . . . . . . . 1.4 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . University of Oxford b University of Toronto c Ought. Correspondence to <EMAIL> 2 Project 2: IDA for Neural Program Interpretation 12 2.1 Motivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12 2.2 Project Directions . . . . . . . . . . . . . . . . . . . . . . . . . . 13 2.3 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 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 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . Background on IDA 0.1 What is IDA?
 Figure 1 : 1 Figure 1: Diagram showing the first few Amplification and Distillation steps in IDA training. First (far left), M is trained by supervised learning on each task x ∈ T0.Second, M solves each task x ∈ T1 by decomposing into tasks {xa, x b } ∈ T0 and solving these directly. This eventually produces a dataset of solved tasks x ∈ T1, which M is trained to solve directly by supervised learning (leftmost "distill" step). This process is then repeated.
 Figure 2 : 2 Figure 2: A multiplication problem could eventually be solved by decomposing all the way down to simple multiplications (without any distillation). Only the first two levels of decomposition are shown here.
 1. 3 . 1 31 Task and DatasetSaxton et al. [24]  introduce a dataset of high-school level mathematics problems in natural language. The problems cover arithmetic, algebra, differentiation, 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
			 More precisely: it is infeasible to provide large numbers of demonstrations or sufficiently dense reward signals for methods like imitation learning or RL to work well.
			 We 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 M to solve them. We could also train distinct models.
			 The Multiplication example has not been implemented, but a version of the Shortest Path example is implemented in [8] .
			 Christiano et al. use a slightly more complicated decomposition. 5 https://ought.org
			 If 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.


Abstract: Reward Machines (RMs) provide a structured, automata-based representation of a reward function that enables a Reinforcement Learning (RL) agent to decompose an RL problem into structured subproblems that can be efficiently learned via off-policy learning. Here we show that RMs can be learned from experience, instead of being specified by the user, and that the resulting problem decomposition can be used to effectively solve partially observable RL problems. We pose the task of learning RMs as a discrete optimization problem where the objective is to find an RM that decomposes the problem into a set of subproblems such that the combination of their optimal memoryless policies is an optimal policy for the original problem. We show the effectiveness of this approach on three partially observable domains, where it significantly outperforms A3C, PPO, and ACER, and discuss its advantages, limitations, and broader potential.  1

Rodrigo Toro Icarte
Ethan Waldie
Toryn Q Klassen
Richard Valenzano
Margarita P Castro
Sheila A Mcilraith
Rodrigo Toro

Learning Reward Machines for Partially Observable Reinforcement Learning

Introduction The use of neural networks for function approximation has led to many recent advances in Reinforcement Learning (RL). Such deep RL methods have allowed agents to learn effective policies in many complex environment including board games  [30] , video games  [23] , and robotic systems  [2] . However, RL methods (including deep RL methods) often struggle when the environment is partially observable. This is because agents in such environments usually require some form of memory to learn optimal behaviour  [31] . Recent approaches for giving memory to an RL agent either rely on recurrent neural networks  [24, 15, 37, 29]  or memory-augmented neural networks  [25, 18] . In this work, we show that Reward Machines (RMs) are another useful tool for providing memory in a partially observable environment. RMs were originally conceived to provide a structured, automatabased representation of a reward function  [33, 4, 14, 39] . Exposed structure can be exploited by the Q-Learning for Reward Machines (QRM) algorithm  [33] , which simultaneously learns a separate policy for each state in the RM. QRM has been shown to outperform standard and hierarchical deep RL over a variety of discrete and continuous domains. However, QRM was only defined for fully observable environments. Furthermore, the RMs were handcrafted. In this paper, we propose a method for learning an RM directly from experience in a partially observable environment, in a manner that allows the RM to serve as memory for an RL algorithm. A requirement is that the RM learning method be given a finite set of detectors for properties that serve as the vocabulary for the RM. We characterize an objective for RM learning that allows us to formulate the task as a discrete optimization problem and propose an efficient local search approach to solve it. By simultaneously learning an RM and a policy for the environment, we are able to significantly outperform several deep RL baselines that use recurrent neural networks as memory in three partially observable domains. We also extend QRM to the case of partial observability where we see further gains when combined with our RM learning method.
 Preliminaries RL agents learn policies from experience. When the problem is fully-observable, the underlying environment model is typically assumed to be a Markov Decision Process (MDP). An MDP is a tuple M = S, A, r, p, γ , where S is a finite set of states, A is a finite set of actions, r : S × A → R is the reward function, p(s, a, s ) is the transition probability distribution, and γ is the discount factor. The agent starts not knowing what r or p are. At every time step t, the agent observes the current state s t ∈ S and executes an action a t ∈ A following a policy π(a t |s t ). As a result, the state s t changes to s t+1 ∼ p(s t+1 |s t , a t ) and the agent receives a reward signal r(s t , a t ). The goal is to learn the optimal policy π * , which maximizes the future expected discounted reward for every state in S  [32] .  [38]  is a well-known RL algorithm that uses samples of experience of the form (s t , a t , r t , s t+1 ) to estimate the optimal q-function q * (s, a). Here, q * (s, a) is the expected return of selecting action a in state s and following an optimal policy π * . Deep RL methods like DQN  [23]  and DDQN  [35]  represent the q-function as qθ (s, a), where qθ is a neural network whose inputs are features of the state and action, and whose weights θ are updated using stochastic gradient descent.
 Q-learning In partially observable problems, the underlying environment model is typically assumed to be a Partially Observable Markov Decision Process (POMDP). A POMDP is a tuple P O = S, O, A, r, p, ω, γ , where S, A, r, p, and γ are defined as in an MDP, O is a finite set of observations, and ω(s, o) is the observation probability distribution. At every time step t, the agent is in exactly one state s t ∈ S, executes an action a t ∈ A, receives reward r t = r(s t , a t ), and moves to state s t+1 according to p(s t , a t , s t+1 ). However, the agent does not observe s t+1 , but only receives an observation o t+1 ∈ O. This observation provides the agent a clue about what the state s t+1 ∈ S is via ω. In particular, ω(s t+1 , o t+1 ) is the probability of observing o t+1 from state s t+1  [5] . RL methods cannot be immediately applied to POMDPs because the transition probabilities and reward function are not necessarily Markovian w.r.t. O (though by definition they are w.r.t. S). As such, optimal policies may need to consider the complete history o 0 , a 0 , . . . , a t−1 , o t of observations and actions when selecting the next action. Several partially observable RL methods use a recurrent neural network to compactly represent the history, and then use a policy gradient method to train it. However, when we do have access to a full POMDP model P O , then the history can be summarized into a belief state. A belief state is a probability distribution b t : S → [0, 1] over S, such that b t (s) is the probability that the agent is in state s ∈ S given the history up to time t. The initial belief state is computed using the initial observation o 0 : b 0 (s) ∝ ω(s, o 0 ) for all s ∈ S. The belief state b t+1 is then determined from the previous belief state b t , the executed action a t , and the resulting observation o t+1 as b t+1 (s ) ∝ ω(s , o t+1 ) s∈S p(s, a t , s )b t (s) for all s ∈ S. Since the state transitions and reward function are Markovian w.r.t. b t , the set of all belief states B can be used to construct the belief MDP M B . Optimal policies for M B are also optimal for the POMDP  [5] .
 Reward Machines for Partially Observable Environments In this section, we define RMs for the case of partial observability. We use the following problem as a running example to help explain various concepts. Example 3.1 (The cookie domain). The cookie domain, shown in Figure  1a , has three rooms connected by a hallway. The agent (purple triangle) can move in the four cardinal directions. There is a button in the yellow room that, when pressed, causes a cookie to randomly appear in the red or blue room. The agent receives a reward of +1 for reaching (and thus eating) the cookie and may then go and press the button again. Pressing the button before reaching a cookie will move it to a random location. There is no cookie at the beginning of the episode. This is a partially observable environment since the agent can only see what it is in the room that it currently occupies.  RMs are finite state machines that are used to encode a reward function  [33] . In the case of partial observability, they are defined over a set of propositional symbols P that correspond to a set of high-level features that the agent can detect using a labelling function L : O ∅ × A ∅ × O → 2 P where (for any set X) X ∅ X ∪ {∅}. L assigns truth values to symbols in P given an environment experience e = (o, a, o ) where o is the observation seen after executing action a when observing o. We use L(∅, ∅, o) to assign truth values to the initial observation. We call a truth value assignment of P an abstract observation because it provides a high-level view of the low-level environment observations via the labelling function L. A formal definition of an RM follows: Definition 3.1 (reward machine). Given a set of propositional symbols P, a Reward Machine is a tuple R P = U, u 0 , δ u , δ r where U is a finite set of states, u 0 ∈ U is an initial state, δ u is the state-transition function, δ u : U ×2 P → U , and δ r is the reward-transition function, δ r : U ×2 P → R. RMs decompose problems into a set of high-level states U and define transitions using if-like conditions defined by δ u . These conditions are over a set of binary properties P that the agent can detect using L. For example, in the cookie domain, P = { , , , , , , }. These properties are true (i.e., part of an experience label according to L) in the following situations: , , , or is true if the agent ends the experience in a room of that color; is true if the agent ends the experience in the same room as a cookie; is true if the agent pushed the button with its last action; and is true if the agent ate a cookie with its last action (by moving onto the space where the cookie was). Figure  2  shows three possible RMs for the cookie domain. They all define the same reward signal (1 for eating a cookie and 0 otherwise) but differ in their states and transitions. As a result, they differ with respect to the amount of information about the current domain state that can be inferred from the current RM state, as we will see below. Each RM starts in the initial state u 0 . Edge labels in the figures provide a visual representation of the functions δ u and δ r . For example, label , 1 between state u 2 and u 0 in Figure  2b  represents δ u (u 2 , { , }) = u 0 and δ r (u 2 , { , }) = 1. Intuitively, this means that if the RM is in state u 2 and the agent's experience ended in room immediately after eating the cookie , then the agent will receive a reward of 1 and the RM will transition to u 0 . Notice that any properties not listed in the label are false (e.g. must be false to take the transition labelled , 1 ). We also use multiple labels separated by a semicolon (e.g., " , 0 ; , 0 ") to describe different conditions for transitioning between the RM states, each with their own associated reward. The label o/w, r ("o/w" for "otherwise") on an edge from u i to u j means that that transition will be made (and reward r received) if none of the other transitions from u i can be taken. Let us illustrate the behaviour of an RM using the one shown in Figure  2c . The RM will stay in u 0 until the agent presses the button (causing a cookie to appear), whereupon the RM moves to u 1 . From u 1 the RM may move to u 2 or u 3 depending on whether the agent finds a cookie when it enters another room. It is also possible to associate meaning with being in RM states: u 0 means that there is no cookie available, u 1 means that there is a cookie in some room (either blue or red), etc. When learning a policy for a given RM, one simple technique is to learn a policy π(o, u) that considers the current observation o ∈ O and the current RM state u ∈ U . Interestingly, a partially observable problem might be non-Markovian over O, but Markovian over O × U for some RM R P . This is the case for the cookie domain with the RM from Figure  2c , for example. Q-Learning for RMs (QRM) is another way to learn a policy by exploiting a given RM  [33] . QRM learns one q-function qu (i.e., policy) per RM state u ∈ U . Then, given any sample experience, the RM can be used to emulate how much reward would have been received had the RM been in any one of its states. Formally, experience e = (o, a, o ) can be transformed into a valid experience ( o, u , a, o , u , r) used for updating qu for each u ∈ U , where u = δ u (u, L(e)) and r = δ r (u, L(e)). Hence, any off-policy learning method can take advantage of these "synthetically" generated experiences to update all subpolicies simultaneously. When tabular q-learning is used, QRM is guaranteed to converge to an optimal policy on fullyobservable problems  [33] . However, in a partially observable environment, an experience e might be more or less likely depending on the RM state that the agent was in when the experience was collected. For example, experience e might be possible in one RM state u i but not in RM state u j . Thus, updating the policy for u j using e as QRM does, would introduce an unwanted bias to quj . We will discuss how to (partially) address this problem in §5. u 0 , 1 ; , 1 ; o/w, 0 (a) Naive RM. u 0 u 1 u 2 o/w, 0 o/w, 0 o/w, 0 , 0 , 0 ; , 0 , 1 , 1 (b) "Optimal" RM. u 0 u 1 u 2 u 3 o/w, 0 o/w, 0 o/w, 0 o/w, 0 , 0 , 0 ; , 0 , 0 ; , 0 , 1 , 1 , 0 , 0 (c) Perfect RM.
 Learning Reward Machines from Traces Our overall idea is to search for an RM that can be used as external memory by an agent for a given task. As input, our method will only take a set of high-level propositional symbols P, and a labelling function L that can detect them. Then, the key question is what properties should such an RM have. Three proposals naturally emerge from the literature. The first comes from the work on learning Finite State Machines (FSMs)  [3, 40, 10] , which suggests learning the smallest RM that correctly mimics the external reward signal given by the environment, as in Giantamidis and Tripakis' method for learning Moore Machines  [10] . Unfortunately, such approaches would learn RMs of limited utility, like the one in Figure  2a . This naive RM correctly predicts reward in the cookie domain (i.e., +1 for eating a cookie , zero otherwise) but provides no memory in support of solving the task. The second proposal comes from the literature on learning Finite State Controllers (FSC)  [22]  and on model-free RL methods  [32] . This work suggests looking for the RM whose optimal policy receives the most reward. For instance, the RM from Figure  2b  is "optimal" in this sense. It decomposes the problem into three states. The optimal policy for u 0 goes directly to press the button, the optimal policy for u 1 goes to the blue room and eats the cookie if present, and the optimal policy for u 2 goes to the red room and eats the cookie. Together, these three policies give rise to an optimal policy for the complete problem. This is a desirable property for RMs, but requires computing optimal policies in order to compare the relative quality of RMs, which seems prohibitively expensive. However, we believe that finding ways to efficiently learn "optimal" RMs is a promising future work direction. Finally, the third proposal comes from the literature on Predictive State Representations (PSR)  [20] , Deterministic Markov Models (DMMs)  [21] , and model-based RL  [16] . These works suggest learning the RM that remembers sufficient information about the history to make accurate Markovian predictions about the next observation. For instance, the cookie domain RM shown in Figure  2c  is perfect w.r.t. this criterion. Intuitively, every transition in the cookie environment is already Markovian except for transitioning from one room to another. Depending on different factors, when entering to the red room there could be a cookie there (or not). The perfect RM is able to encode such information using 4 states: when at u 0 the agent knows that there is no cookie, at u 1 the agent knows that there is a cookie in the blue or the red room, at u 2 the agent knows that there is a cookie in the red room, and at u 3 the agent knows that there is a cookie in the blue room. Since keeping track of more information will not result in better predictions, this RM is perfect. Below, we develop a theory about perfect RMs and describe an approach to learn them.
 Perfect Reward Machines: Formal Definition and Properties The key insight behind perfect RMs is to use their states U and transitions δ u to keep track of relevant past information such that the partially observable environment P O becomes Markovian w.r.t. O × U . Definition 4.1 (perfect reward machine). An RM R P = U, u 0 , δ u , δ r is considered perfect for a POMDP P O = S, O, A, r, p, ω, γ with respect to a labelling function L if and only if for every trace o 0 , a 0 , . . . , o t , a t generated by any policy over P O , the following holds: Pr(o t+1 , r t |o 0 , a 0 , . . . , o t , a t ) = Pr(o t+1 , r t |o t , x t , a t ) (1) where x 0 = u 0 and x t = δ u (x t−1 , L(o t−1 , a t−1 , o t )) . Two interesting properties follow from Definition 4.1. First, if the set of belief states B for the POMDP P O is finite, then there exists a perfect RM for P O with respect to some L. Second, the optimal policies for perfect RMs are also optimal for the POMDP (see supplementary material §12). Instead, we propose an alternative that focuses on a necessary condition for a perfect RM: the RM must predict what is possible and impossible in the environment at the abstract level. For example, it is impossible to be at u 3 in the RM from Figure  2c  and make the abstract observation { , }, because the RM reaches u 3 only if the cookie was seen in the blue room or not to be in the red room. This idea is formalized in the optimization model LRM. Let T = {T 0 , . . . , T n } be a set of traces, where each trace T i is a sequence of observations, actions, and rewards: T i = (o i,0 , a i,0 , r i,0 , . . . , a i,ti−1 , r i,ti−1 , o i,ti ). (2) We now look for an RM U, u 0 , δ u , δ r that can be used to predict L(e  |U | ≤ umax (4) xi,t ∈ U ∀i ∈ I, t ∈ Ti ∪ {ti} (5) xi,0 = u0 ∀i ∈ I (6) xi,t+1 = δu(xi,t, L(ei,t+1)) ∀i ∈ I, t ∈ Ti (7) N u,l ⊆ 2 2 P ∀u ∈ U, l ∈ 2 P (8) L(ei,t+1) ∈ N x i,t ,L(e i,t ) ∀i ∈ I, t ∈ Ti (9) Constraints (  3 ) and (  4 ) ensure that we find a well-formed RM over P with at most u max states. Constraint (  5 ),  (6) , and (7) ensure that x i,t is equal to the current state of the RM, starting from u 0 and following δ u . Constraint (  8 ) and (  9 ) ensure that the sets N u,l contain every L(e i,t+1 ) that has been seen right after l and u in T . The objective function comes from maximizing the log-likelihood for predicting L(e i,t+1 ) using a uniform distribution over all the possible options given by N u,l . A key property of this formulation is that any perfect RM is optimal w.r.t. the objective function in LRM when the number of traces tends to infinity (see supplementary material §12): Theorem 4.3. When the set of training traces (and their lengths) tends to infinity and is collected by a policy such that π(a|o) > for all o ∈ O and a ∈ A, any perfect RM with respect to L and at most u max states will be an optimal solution to the formulation LRM. Finally, note that the definition of a perfect RM does not impose conditions over the rewards associated with the RM (i.e., δ r ). This is why δ r is a free variable in the model LRM. However, we still expect δ r to model the external reward signals given by the environment. To do so, we estimate δ r (u, l) using its empirical expectation over T (as commonly done when constructing belief MDPs  [5] ).
 Searching for a Perfect Reward Machine Using Tabu Search We now describe the specific optimization technique used to solve LRM. We experimented with many discrete optimization approaches-including mixed integer programming  [6] , Benders decomposition  [8] , evolutionary algorithms  [17] , among others-and found local search algorithms  [1]  to be the most effective at finding high quality RMs given short time limits. In particular, we use Tabu search  [11] , a simple and versatile local search procedure with convergence guarantees and many successful applications in the literature  [36] . We also include our unsuccessful mixed integer linear programming model for LRM in the supplementary material §10. In the context of our work, Tabu search starts from a random RM and, on each iteration it evaluates all "neighbouring" RMs. We define the neighbourhood of an RM as the set of RMs that differ by exactly one transition (i.e., removing/adding a transition, or changing its value) and evaluate RMs using the objective function of LRM. When all neighbouring RMs are evaluated, the algorithm chooses the one with lowest values and sets it as the current RM. To avoid local minima, Tabu search maintains a Tabu list of all the RMs that were previously used as the current RM. Then, RMs in the Tabu list are pruned when examining the neighbourhood of the current RM.
 Simultaneously Learning a Reward Machine and a Policy We now describe our overall approach to simultaneously finding an RM and exploiting that RM to learn a policy. The complete pseudo-code can be found in the supplementary material (Algorithm 1). Our approach starts by collecting a training set of traces T generated by a random policy during t w "warmup" steps. This set of traces is used to find an initial RM R using Tabu search. The algorithm then initializes policy π, sets the RM state to the initial state u 0 , and sets the current label l to the initial abstract observation L(∅, ∅, o). The standard RL learning loop is then followed: an action a is selected following π(o, u) where u is the current RM state, and the agent receives the next observation o and the immediate reward r. The RM state is then updated to u = δ u (u, L(o, a, o )) and the last experience ( o, u , a, r, o , u ) is used by any RL method of choice to update π. Note that in an episodic task, the environment and RM are reset whenever a terminal state is reached. If on any step, there is evidence that the current RM might not be the best one, our approach will attempt to find a new one. Recall that the RM R was selected using the cardinality of its prediction sets N (LRM). Hence, if the current abstract observation l is not in N u,l , adding the current trace to T will increase the size of N u,l for R. As such, the cost of R will increase and it may no longer be the best RM. Thus, if l ∈ N u,l , we add the current trace to T and search for a new RM. Recall that we use Tabu search, though any discrete optimization method could be applied. Our method only uses the new RM if its cost is lower than R's. If the RM is updated, a new policy is learned from scratch. Given the current RM, we can use any RL algorithm to learn a policy π(o, u), by treating the combination of o and u as the current state. If the RM is perfect, then the optimal policy π * (o, u) will also be optimal for the original POMDP (as stated in Theorem 4.2). However, to exploit the problem structure exposed by the RM, we can use the QRM algorithm. As explained in §3, standard QRM under partial observability can introduce a bias because an experience e = (o, a, o ) might be more or less likely depending on the RM state that the agent was in when the experience was collected. We partially address this issue by updating qu using (o, a, o ) if and only if L(o, a, o ) ∈ N u,l , where l was the current abstract observation that generated the experience (o, a, o ). Hence, we do not transfer experiences from u i to u j if the current RM does not believe that (o, a, o ) is possible in u j . For example, consider the cookie domain and the perfect RM from Figure  2c . If some experience consists of entering to the red room and seeing a cookie, then this experience will not be used by states u 0 and u 3 as it is impossible to observe a cookie at the red room from those states. Note that adding this rule may work in many cases, but it will not address the problem in all environments (more discussion in §7). We consider addressing this problem as an interesting area for future work.
 Experimental Evaluation We tested our approach on three partially observable grid domains (Figure  1 ). The agent can move in the four cardinal directions and can only see what is in the current room. These are stochastic domains where the outcome of an action randomly changes with a 5% probability. The first environment is the cookie domain (Figure  1a ) described in §3. Each episode is 5, 000 steps long, during which the agent should attempt to get as many cookies as possible. The second environment is the symbol domain (Figure  1b ). It has three symbols ♣, ♠, and in the red and blue rooms. One symbol from {♣, ♠, } and possibly a right or left arrow are randomly placed at the yellow room. Intuitively, that symbol and arrow tell the agent where to go, e.g., ♣ and → tell the agent to go to ♣ in the east room. If there is no arrow, the agent can go to the target symbol in either room. An episode ends when the agent reaches any symbol in the red or blue room, at which point it receives a reward of +1 if it reached the correct symbol and −1 otherwise. The third environment is the 2-keys domain (Figure  1c ). The agent receives a reward of +1 when it reaches the coffee (in the yellow room). To do so, it must open the two doors (shown in brown). Each door requires a different key to open it, and the agent can only carry one key at a time. Initially, the two keys are randomly located in either the blue room, the red room, or split between them. We tested two versions of our Learned Reward Machine (LRM) approach: LRM+DDQN and LRM+DQRM. Both learn an RM from experience as described in §4.2, but LRM+DDQN learns a policy using DDQN  [35]  while LRM+DQRM uses the modified version of QRM described in §5. In all domains, we used u max = 10, t w = 200, 000, an epsilon greedy policy with = 0.1, and a discount factor γ = 0.9. The size of the Tabu list and the number of steps that the Tabu search performs before returning the best RM found is 100. We compared against 4 baselines: DDQN  [35] , A3C  [24] , ACER  [37] , and PPO  [29]  using the OpenAI baseline implementations  [12] . DDQN uses the concatenation of the last 10 observations as input which gives DDQN a limited memory to better handle the domains. A3C, ACER, and PPO use an LSTM to summarize the history. Note that the output of the labelling function was also given to the baselines. Details on the hyperparameters and networks can be found in the supplementary material §13. Figure  3  shows the total cumulative rewards that each approach gets every 10, 000 training steps and compares it to the optimal policy. For the LRM algorithms, the figure shows the median performance over 30 runs per domain, and percentile 25 to 75 in the shadowed area. For the DDQN baseline, we show the maximum performance seen for each time period over 5 runs per problem. Similarly, we also show the maximum performance over the 30 runs of A3C, ACER, and PPO per period. All the baselines outperformed a random policy, but none make much progress on any of the domains. Furthermore, LRM approaches largely outperform all the baselines, reaching close-to-optimal policies in the cookie and symbol domain. We also note that LRM+DQRM learns faster than LRM+DDQN, but is more unstable. In particular, LRM+DQRM converged to a considerably better policy than LRM+DDQN in the 2-keys domain. We believe this is due to QRM's experience sharing mechanism that allows for propagating sparse reward backwards faster (see supplementary material §13.3). A key factor in the strong performance of the LRM approaches is that Tabu search finds high-quality RMs in less than 100 local search steps (Figure  5 , supplementary material). In fact, our results show that Tabu search finds perfect RMs in most runs, in particular when tested over the symbol domain. 7 Discussion, Limitations, and Broader Potential Solving partially observable RL problems is challenging and LRM was able to solve three problems that were conceptually simple but presented a major challenge to A3C, ACER, and PPO with LSTM-based memories. A key idea behind these results was to optimize over a necessary condition for perfect RMs. This objective favors RMs that are able to predict possible and impossible future observations at the abstract level given by the labelling function L. In this section, we discuss the advantages and current limitations of such an approach. u 0 u 1 , 1 ; o/w, 0 , 1 ; o/w, 0 , 0 , 0 We begin by considering the performance of Tabu search in our domains. Given a training set composed of one million transitions, a simple Python implementation of Tabu search takes less than 2.5 minutes to learn an RM across all our environments, when using 62 workers on a Threadripper 2990WX processor. Note that Tabu search's main bottleneck is evaluating the neighbourhood around the current RM solution. As the size of the neighbourhood depends on the size of the set of propositional symbols P, exhaustively evaluating the neighbourhood may sometimes become impractical. To handle such problem, it will be necessary to import ideas from the Large Neighborhood Search literature  [27] . Regarding limitations, learning the RM at the abstract level is efficient but requires ignoring (possibly relevant) low-level information. For instance, Figure  4  shows an adversarial example for LRM. The agent receives reward for eating the cookie ( ). There is an external force pulling the agent down-i.e., the outcome of the "move-up" action is actually a downward movement with high probability. There is a button ( ) that the agent can press to turn off (or back on) the external force. Hence, the optimal policy is to press the button and then eat the cookie. Given P = { , }, a perfect RM for this environment is fairly simple (see Figure  4 ) but LRM might not find it. The reason is that pressing the button changes the low-level probabilities in the environment but does not change what is possible or impossible at the abstract level. In other words, while the LRM objective optimizes over necessary conditions for finding a perfect RM, those conditions are not sufficient to ensure that an optimal solution will be a perfect RM. In addition, if a perfect RM is found, our heuristic approach to share experiences in QRM would not work as intended because the experiences collected when the force is on (at u 0 ) would be used to update the policy for the case where the force is off (at u 1 ). Other current limitations include that it is unclear how to handle noise over the high-level detectors L and how to transfer learning from previously learned policies when a new RM is learned. Finally, defining a set of proper high-level detectors for a given environment might be a challenge to deploying LRM. Hence, looking for ways to automate that step is an important direction for future work.
 Related Work State-of-the-art approaches to partially observable RL use Recurrent Neural Networks (RNNs) as memory in combination with policy gradient  [24, 37, 29, 15] , or use external neural-based memories  [25, 18, 13] . Other approaches include extensions to Model-Based Bayesian RL to work under partial observability  [28, 7, 9]  and to provide a small binary memory to the agent and a special set of actions to modify it  [26] . While our experiments highlight the merits of our approach w.r.t. RNN-based approaches, we rely on ideas that are largely orthogonal. As such, we believe there is significant potential in mixing these approaches to get the benefit of memory at both the high-and the low-level. The effectiveness of automata-based memory has long been recognized in the POMDP literature  [5] , where the objective is to find policies given a complete specification of the environment. The idea is to encode policies using Finite State Controllers (FSCs) which are FSMs where the transitions are defined in terms of low-level observations from the environment and each state in the FSM is associated with one primitive action. When interacting with the environment, the agent always selects the action associated with the current state in the controller. Meuleau et al.  [22]  adapted this idea to work in the RL setting by exploiting policy gradient to learn policies encoded as FSCs. RMs can be considered as a generalization of FSC as they allow for transitions using conditions over high-level events and associate complete policies (instead of just one primitive action) to each state. This allows our approach to easily leverage existing deep RL methods to learn policies from low-level inputs, such as images-which is not achievable by Meuleau et al.  [22] . That said, further investigating using ideas for learning FSMs  [3, 40, 10]  in learning RMs is a promising direction for future work. Our approach to learn RMs is greatly influenced by Predictive State Representations (PSRs)  [20] . The idea behind PSRs is to find a set of core tests (i.e., sequences of actions and observations) such that if the agent can predict the probabilities of these occurring, given any history H, then those probabilities can be used to compute the probability of any other test given H. The insight is that state representations that are good for predicting the next observation are good for solving partially observable environments. We adapted this idea to the context of RM learning as discussed in §4. While our work was under review, two interesting papers were submitted to arXiv. The first paper, by Xu et al.  [39] , proposes a polynomial time algorithm to learn reward machines in fully observable domains. Their goal is to learn the smallest reward machine that is consistent with the reward function-which makes sense for fully observable domains, but would have limited utility under partial observability (as discussed in §4). The second paper, by Zhang et al.  [41] , proposes to learn a discrete PSR representation of the environment directly from low-level observations and then plan over such representation using tabular Q-learning. This is a promising research direction, with some clear synergies with LRM.
 Concluding Remarks We have presented a method for learning (perfect) Reward Machines in partially observable environments and demonstrated the effectiveness of these learned RMs in tackling partially observable RL problems that are unsolvable by A3C, ACER and PPO. Informed by criteria from the POMDP, FSC, and PSR literature, we proposed a set of RM properties that support tackling RL in partially observable environments. We used these properties to formulate RM learning as a discrete optimization problem. We experimented with several optimization methods, finding Tabu search to be the most effective. We then combined this RM learning with policy learning for partially observable RL problems. Our combined approach outperformed a set of strong LSTM-based approaches on different domains. We believe this work represents an important building block for creating RL agents that can solve cognitively challenging partially observable tasks. Not only did our approach solve problems that were unsolvable by A3C, ACER and PPO, but it did so in a relatively small number of training steps. RM learning provided the agent with memory, but more importantly the combination of RM learning and policy learning provided it with discrete reasoning capabilities that operated at a higher level of abstraction, while leveraging deep RL's ability to learn policies from low-level inputs. This work leaves open many interesting questions relating to abstraction, observability, and properties of the language over which RMs are constructed. We believe that addressing these questions, among many others, will push the boundary of partially observable RL problems that can be solved.
 Algorithm for Simultaneously Learning Reward Machines and a Policy Algorithm 1 shows our overall approach to simultaneously learning an RM and exploiting that RM to learn a policy. The algorithm inputs are the set of propositional symbols P, the labelling function L, a maximum on the number of RM states u max , and the number of "warmup" steps t w . Our approach starts by collecting a training set of traces T generated by a random policy during t w steps (Line 2). This set of traces is used to find an initial RM R using Tabu search (Line 3). If later traces show that R is incorrect, our method will then find a new RM learned using the additional traces. Lines 4 and 5 initialize the environment and the policy π, and set variables x and l to the initial RM state u 0 and initial abstract observation L(∅, ∅, o), respectively. Lines 7-19 are the main loop of our approach. Lines 7-10 are part of the standard RL loop: the agent executes an action a selected following π(o, x) and receives the next observation o , the immediate reward r, and a boolean variable done indicating if the episode has terminated. Then, the state in the RM x is updated and the policy π is improved using the last experience ( o, x , a, r, o , x , done). Note that when done is true, the environment and RM are reset (Lines 17-18). Lines 11-16 involve relearning the RM when there is evidence that the current RM might not be the best one. Recall that the RM R was selected using the cardinality of its prediction sets N (see the description of LRM). Hence, if the current abstract observation l is not in N x,l , then adding the current trace to T will increase the size of N x,l for R. As such, the cost of R will increase and it may no longer be the best RM. Thus, if l ∈ N x,l , we add the current trace to T and use Tabu search to find a new RM. Note, our method only uses the new RM if its cost is lower than that of R (Lines 14-16). However, when the RM is updated, a new policy is learned from scratch (Line 16). T ← T ∪ get_current_trace() Proof sketch. If the reachable belief space B is finite, we can construct an RM that keeps track of the current belief state using one RM state per belief state and emulating their progression using δ u , and one propositional symbol for every action-observation pair. Thus, the current belief state b t can be inferred from the last observation, last action, and the current RM state. As such, the equality from Definition 4.1 holds. Two interesting properties follow from the definition of a perfect RM. First, if the set of belief states B for the POMDP P O is finite, then there exists a perfect RM for P O with respect to some L. Second, the optimal policies for perfect RMs are also optimal for the POMDP. Theorem 12.2. Let R P be a perfect RM for a POMDP P O w.r.t. a labelling function L, then any optimal policy for R P w.r.t. the environmental reward is also optimal for P O . Proof sketch. As the next observation and immediate reward probabilities can be predicted from O × U × A, an optimal policy over O × U must also be optimal over P O . A key property of this formulation is that any perfect RM is optimal with respect to the objective function in LRM when the number of traces tends to infinity: Theorem 12.3. When the set of training traces (and their lengths) tends to infinity and is collected by a policy such that π(a|o) > for all o ∈ O and a ∈ A, and some constant > 0, then any perfect RM with respect to L and at most u max states will be an optimal solution to the formulation given in LRM. Proof sketch. In the limit, l ∈ N u,l if and only if the probability of observing l after executing an action from the RM state u while observing l is non-zero. In particular, for all i ∈ I and t ∈ T , the cardinality of N xi,t,L(ei,t) will be minimal for a perfect RM. This follows from the fact that perfect RMs make perfect predictions for the next observation o given o, u, and a. Therefore, as we minimize the sum over log(|N xi,t,L(ei,t) |), the objective value for a perfect RM must be minimal.
 Experimental Evaluation 13.1 Experimental Details For LRM+DDQN and LRM+DQRM, the neural network used has 5 fully connected layers with 64 neurons per layer. On every step, we trained the network using 32 sampled experiences from a replay buffer of size 100,000 using the Adam optimizer  [19]  and a learning rate of 5e-5. The target networks were updated every 100 steps. DDQN  [35]  uses the same parameters and network architecture as LRM+DDQN, but its input is the concatenation of the last 10 observations, as commonly done by Atari playing agents. This gives DDQN a limited memory to better handle partially observable domains. In contrast, A3C, ACER, and PPO use an LSTM to summarize the history. We also followed the same testing methodology that was used in their original publications. We ran each approach at least 30 times per domain, and on every run, we randomly selected the number of hidden neurons for the LSTM from {64, 128, 256, 512} and a learning rate from (1e-3, 1e-5). We also sampled δ from {0, 1, 2} for ACER and the clip range from (0.1, 0.3) for PPO. Other parameters were fixed to their default values. While interacting with the environment, the agents were given a "top-down" view of the world represented as a set of binary matrices. One matrix had a 1 in the current location of the agent, one had a 1 in only those locations that are currently observable, and the remaining matrices each corresponded to an object in the environment and had a 1 at only those locations that were both currently observable and contained that object (i.e., locations in other rooms are "blacked out"). The agent also had access to features indicating if they were carrying a key, which colour room they were in, and the current status (i.e., occurring or not occurring) of the events detected by the labelling function.
 Tabu Search Figure  5  evaluates the quality of the RMs found by Tabu search by comparing it the perfect RM. In each plot, a dot compares the cost of a handcrafted perfect RM with that of an RM R that was found by Tabu search while running our LRM approaches, where the costs are evaluated relative to the training set used to find R. Being on or under the diagonal line (as in most of the points in the figure) means that Tabu search is finding RMs whose values are at least as good as the handcrafted RM. Hence, Tabu search is either finding perfect RMs or discovering that our training set is incomplete and our agent will eventually fill those gaps.
 Cookie domain Symbol domain 2-keys domains Cost Perfect RM Cost Learned RM Cost Perfect RM Cost Learned RM Cost Perfect RM Cost Learned RM
 DDQN vs DQRM: Exploration Heatmaps and Learned Trajectories As shown in §6 of the paper, LRM+DQRM tends to learn faster than LRM+DDQN and largely outperforms LRM+DDQN in the 2-keys domain. Towards better understating these results, we ran the following experiment. For the 2-keys domain (and identical random seed), we learn policies using DDQN and DQRM over the same handcrafted perfect RM from Figure  6 . RM state u 6 is fairly simple: when the agent is carrying a key and there is only one closed door, go and get the coffee. As expected, DQRM seems to learn an accurate estimation of that q-function using very limited interactions within the coffee room. In contrast, DDQN uses one big q-network to model the complete policy. Receiving reward by getting the coffee pushes the q-network estimation to believe that a high reward can be collected from the coffee room (even if the doors are closed). Hence, the agent spent a considerable amount of time hitting the second door without having a key. Second, DQRM shares experience over all the q-networks. This allows for using the experiences collected while being at early stages of the task (e.g., states u 0 , u 1 , u 2 , and u 3 ) to update policies for later stages (e.g., states u 4 , u 5 , and u 6 ). In particular, all the experience needed to learn to navigate from one room to another is shared. Therefore, while DDQN would depend on its network to avoid relearning how to navigate between rooms when the RM state changes, DQRM enforces such transfer by sharing experiences whenever it is appropriate. This might explain why DDQN spends considerably more time in the hallway than DQRM in heatmaps 50-100K and 100-150K.  Note that the exploratory trend from 150-500K follows a clear pattern. On one hand, the DQRM agent seems to have a good idea of how to solve the task and, therefore, it spends most of its time on the hallway (solving the tasks requires passing through the hallway at least 4 times). On the other hand, the DDQN agent is getting stuck exploring subregions of the map. Finally, we inspected the trajectories of each agent when solving the task after 1 million training steps. Figures 9 shows the DDQN agent and Figure  10  shows the DQRM agent. Both agents solved the task, but DQRM solved it faster (83 steps vs 102 steps). In solving this problem, the main difficulty for the DDQN agent is in reacting to entering the South room and discovering that the keys are not there. Its reaction is to enter and leave the empty room many times before going to the North room. In the case of DQRM, the agent goes directly to the North room after observing that the South room is empty, but then it enters and leaves the North room a few times before collecting the keys. After collecting the first key, both agents solved the rest of the task almost optimally.    Figure 1 : 1 Figure 1: Partially observable environments where the agent can only see what is in the current room.
 Figure 2 : 2 Figure 2: Three possible Reward Machines for the Cookie domain.
 Figure 3 : 3 Figure 3: Total reward collected every 10, 000 training steps.
 Figure 4 : 4 Figure 4: The gravity domain.
 Algorithm 1 1 Learning an RM and a Policy 1: Input: P, L, A, u max , t w 2: T ← collect_traces(t w ) 3: R, N ← learn_rm(P, L, T , u max ) 4: o, x, l ← env_get_initial_state(), u 0 , L(∅, ∅, o) 5: π ← initialize_policy() 6: for t = 1 to t train do7: a ← select_action(π, o, x) 8: o , r, done ← env_execute_action(a) 9: x , l ← δ u (x, L(o, a, o )), L(o, a, o ) 10: π ← improve(π, o, x, l, a, r, o , x , l , done, N ) 11: if l ∈ N x,l then 12:
 Figure 5 : 5 Figure 5: Cost comparison between perfect RM and RM found by Tabu search.
 steps 300-350K steps 350-400K steps 400-450K steps 450-500K steps DQRM DQRM DQRM DQRM DQRM 250-300K steps 300-350K steps 350-400K steps 400-450K steps 450-500K steps
 Figure 8 : 8 Figure 8: Exploration heatmaps for DDQN and DQRM over the 2-keys domain given a perfect RM.
 (a) Collecting the first key. (b) Opening the first door. (c) Collecting the other key. (d) Opening the second door.
 Figure 9 : 9 Figure 9: The learned trace after 1 million training steps by DDQN given a perfect RM, divided into four stages.
 Figure 10 : 10 Figure 10: The learned trace after 1 million training steps by QRM given a perfect RM, divided into four stages.
 i,t+1 ) from L(e i,t ) and the current RM state x i,t , where e i,t+1 is the experience (o i,t , a i,t , o i,t+1 ) and e i,0 is (∅, ∅, o i,0 ) by definition. The model parameters are the set of traces T , the set of propositional symbols P, the labelling function L, and a maximum number of states in the RM u max . The model also uses the sets I = {0 . . . n} and T i = {0 . . . t i − 1}, where I contains the index of the traces and T i their time steps. The model has two auxiliary variables x i,t and N u,l . Variable x i,t ∈ U represents the state of the RM after observing trace T i up to time t. Variable N u,l ⊆ 2 2 P is the set of all the next abstract observations seen from the RM state u and the abstract observations l at some point in T . In other words, l ∈ N u,l iff u = x i,t , l = L(e i,t ), and l = L(e i,t+1 ) for some trace T i and time t. minimize U,u 0 ,δu,δr i∈I t∈T i log(|N x i,t ,L(e i,t ) |) (LRM) s.t. U, u0, δu, δr ∈ RP (3)
 Given any POMDP P O with a finite reachable belief space, there exist a perfect RM for P O with respect to some labelling function L. 12 Theorems and Proof Sketches Theorem 12.1. 13: 14: 15: 16: 17: R , N ← relearn_rm(R, P, L, T , u max ) if R = R then R, done ← R , true π ← initialize_policy() end if 18: 19: 20: 21: end if if done then o , x , l ← env_get_initial_state(), u 0 , L(∅, ∅, o) end if 22: 23: end for o, x, l ← o , x , l 24: return π