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 andreas@ought.org 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: 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.


Responsive Safety in Reinforcement Learning by PID Lagrangian Methods

Introduction Reinforcement learning has solved sequential decision tasks of impressive difficulty by maximizing reward functions through trial and error. Recent examples using deep learning range from robotic locomotion  (Schulman et al., 2015; Gu et al., 2016; Schulman et al., 2017; Levine et al., 2016)  to sophisticated video games  (Mnih et al., 2013; Schulman et al., 2017; OpenAI, 2018; Jaderberg et al., 2019) . While errors during training in these domains come without cost, in some learning scenarios it is important to limit the rates of hazardous outcomes. One example would be wear and tear 1 University of California, Berkeley 2 OpenAI. Correspondence to: Adam Stooke <adam.stooke@berkeley.edu>. Proceedings of the 37 th International Conference on Machine  Learning, Vienna, Austria, PMLR 119, 2020.  Copyright 2020 by the author(s). on a robot's components or its surroundings. It may not be possible to impose such limits by prescribing constraints in the action or state space directly; instead, hazard-avoiding behavior must be learned. For this purpose, we use the well-known framework of the constrained Markov decision process (CMDP)  (Altman, 1999) , which limits the accumulation of a "cost" signal which is analogous to the reward. The optimal policy is one which maximizes the usual return while satisfying the cost constraint. In safe RL the agent must avoid hazards not only at convergence, but also throughout exploration and learning. Lagrangian methods are a classic approach to solving constrained optimization problems. For example, the equalityconstrained problem over the real vector x: min x f (x) s.t. g(x) = 0 (1) is transformed into an unconstrained one by introduction of a dual variable-the Lagrange multiplier, λ-to form the Lagrangian: L(x, λ) = f (x) + λg(x), which is used to find the solution as: (x * , λ * ) = arg max λ min x L(x, λ) (2) Gradient-based algorithms iteratively update the primal and dual variables: −∇ x L(x, λ) = − ∇ x f (x) − λ∇ x g(x) (3) ∇ λ L(x, λ) = g(x) (4) so that λ acts as a learned penalty coefficient in the objective, leading eventually to a constraint-satisfying solution (see e.g.  Bertsekas (2014) ). The Lagrangian multiplier method is readily adapted to the constrained RL setting  (Altman, 1998; Geibel & Wysotzki, 2011)  and has become a popular baseline in deep RL  (Achiam et al., 2017; Chow et al., 2019)  for its simplicity and effectiveness. Although they have been shown to converge to optimal, constraint-satisfying policies  (Tessler et al., 2018; Paternain et al., 2019) , a shortcoming of gradient Lagrangian methods for safe RL is that intermediate iterates often violate constraints. Cost overshoot and oscillations are in fact inherent to the learning dynamics  (Platt & Barr, 1988; Wah et al., 2000) , and we witnessed numerous problematic cases in our own experiments. Figure  1  (left) shows an example from a deep RL setting, where the cost and multiplier values oscillated throughout training. Our key insight in relation to this deficiency is that the traditional Lagrange multiplier update in (4) amounts to integral control on the constraint. The 90-degree phase shift between the curves is characteristic of ill-tuned integral controllers. Our contribution is to expand the scope of possible Lagrange multiplier update rules beyond (4), by interpreting the overall learning algorithm as a dynamical system. Specifically, we employ the next simplest mechanisms, proportional and derivative control, to λ, by adding terms corresponding to derivatives of the constraint function into (4) (derivatives with respect to learning iteration). To our knowledge, this is the first time that an expanded update rule has been considered for a learned Lagrange multiplier. PID control is an appealing enhancement, evidenced by the fact that it is one of the most widely used and studied control techniques  ( Åström & Hägglund, 2006) . The result is a more responsive safety mechanism, as demonstrated in Figure  1  (right), where the cost oscillations have been damped, dramatically reducing violations. Our contributions in this paper are outlined as follows. First, we provide further context through related works and preliminary definitions. In Section 4, we propose modified Lagrangian multiplier methods and analyze their benefits in the learning dynamics. Next, in Section 5, we cast constrained RL as a dynamical system with the Lagrange multiplier as a control input, to which we apply PID control as a new algorithm. In Section 6, we adapt a leading deep RL algorithm, Proximal Policy Optimization (PPO)  (Schulman et al., 2017)  with our methods and achieve state of the art performance in the OpenAI Safety-Gym suite of environments  (Ray et al., 2019) . Finally, in Section 7 we introduce another novel technique that makes tuning easier by providing invariance to the relative numerical scales of rewards and costs, and we demonstrate it in a further set of experiments. Our extensive empirical results show that our algorithms, which are intuitive and simple to implement, improve cost performance and promote hyperparameter robustness in a deep RL setting.
 Related Work Constrained Deep RL. Adaptations of the Lagrange multiplier method to the actor-critic RL setting have been shown to converge to the optimal, constraint-satisfying solution under certain assumptions  (Tessler et al., 2018) . Convergence proofs have relied upon updating the multiplier more slowly than the policy parameters  (Tessler et al., 2018; Paternain et al., 2019) , implying many constraint-violating policy iterations may occur before the penalty comes into full effect. Several recent works have aimed at improving constraint satisfaction in RL over the Lagrangian method, but they tend to incur added complexity.  Achiam et al. (2017)  introduced Constrained Policy Optimization (CPO), a policy search algorithm with near-constraint satisfaction guarantees at every iteration, based on a new bound on the expected returns of two nearby policies. CPO includes a projection step on the policy parameters, which in practice requires a time-consuming backtracking line search. Yet, simple Lagrangian-based algorithms performed as well or better in a recent empirical comparison in Safety Gym  (Ray et al., 2019) . Approaches to safe RL based on Lyapunov functions have been developed in a series of studies  (Chow et al., 2018; , resulting in algorithms that combine a projection step, as in CPO, with action-layer interventions like the safety layer of  Dalal et al. (2018) . Experimentally, this line of work showed mixed performance gains over Lagrangian methods, at a nontrivial cost to implement and without clear guidance for tuning.  developed interior point methods for RL, which augment the objective with logarithmic barrier functions. These methods are shown theoretically to provide suboptimal solutions. Furthermore, they require tuning of the barrier strength and typically assume already feasible iterates, the latter point possibly being problematic for random agent initializations or under noisy cost estimates. Most recently,  Yang et al. (2020)  extended CPO with a two-step projection-based optimization approach. In contrast to these techniques, our method remains nearly as simple to implement and compute as the baseline Lagrangian method. Dynamical Systems View of Optimization. Several recent works have proposed different dynamical systems viewpoints to analyze optimization algorithms, including those often applied to deep learning.  Hu & Lessard (2017)  reinterpreted first-order gradient optimization as a dynamical system; they likened the gradient of the objective, ∇ x f , to the plant, which the controller aims to drive to zero to arrive at the optimal parameters, x * . Basic gradient de-scent then matches the form of integral control (on ∇ x f ). They extend the analogy to momentum-based methods, for example linking Nesterov momentum to PID control with lag compensation. In another example,  An et al. (2018)  interpreted SGD as P-control and momentum methods as PI-control. They introduced a derivative term, based on the change in the gradient, and applied their resulting PID controller to improve optimization of deep convolutional networks. Other recent works bring yet other perspectives from dynamical systems to deep learning and optimization, see for example  (Lessard et al., 2014; Nishihara et al., 2015; Liu & Theodorou, 2019) ). None of these works address constrained RL, however, necessitating our distinct formulation for that problem. Constrained Optimization. Decades' worth of literature have accumulated on Lagrangian methods. But even recent textbooks on the topic  (Bertsekas, 2014; Nocedal & Wright, 2006)  only consider updating the Lagrange multiplier using the value of the constraint function, g(x), and miss ever using its derivatives, ġ(x) or g(x), which we introduce. The modification to the Lagrangian method most similar in effect to our proportional control term (here using ġ(x)) is the quadratic penalty method  (Hestenes (1969) ;  Powell (1969)  see also e.g.  Bertsekas (1976) ), which we compare in Section 4.  Song & Leland (1998)  proposed a controls viewpoint (continuous-time) of optimizing neural networks for constrained problems and arrived at proportional control rules only. Related to our final experiments on reward-scale invariance,  Wah et al. (2000)  developed an adaptive weighting scheme for continuous-time Lagrangian objectives, but it is an intricate procedure which is not straightforwardly applied to safe RL.
 Preliminaries Constrained Reinforcement Learning Constrained Markov Decision Processes (CMDP)  (Altman, 1998)  extend MDPs (see  Sutton & Barto (1998) ) to incorporate constraints into reinforcement learning. A CMDP is the expanded tuple (S, A, R, T, µ, C 0 , C 1 , ..., d 0 , d 1 , ...), with the cost functions C i : S × A × S → R defined by the same form as the reward, and d i : R denoting limits on the costs. For ease of notation, we will only consider a single, all-encompassing cost. The expected sum of discounted rewards over trajectories, τ = (s 0 , a 0 , s 1 , a 1 , ...), induced by the policy π(a|s) is a common performance objective: J(π) = E τ ∼π [ ∞ t=0 γ t R(s t , a t , s t+1 )]. The analogous value function for cost is defined as: J C (π) = E τ ∼π [ ∞ t=0 γ t C(s t , a t , s t+1 )]. The constrained RL problem is to solve for the best feasible policy: π * = arg max π J(π) s.t. J C (π) ≤ d (5) Deep reinforcement learning uses a deep neural network for the policy, π θ = π(•|s; θ) with parameter vector θ, and policy gradient algorithms improve the policy iteratively by gathering experience in the task to estimate the reward objective gradient, ∇ θ J(π θ ). Thus our problem of interest is better expressed as maximizing score at some iterate, π k , while ideally obeying constraints at each iteration: max π J(π k ) s.t. J C (π m ) ≤ d m ∈ {0, 1, ..., k} (6) Practical settings often allow trading reward performance against some constraint violations (e.g. the constraints themselves may include a safety margin). For this purpose we introduce a constraint figure of merit with our experiments.
 Dynamical Systems and Optimal Control Dynamical systems are processes which can be subject to an external influence, or control. A general formulation for discrete-time systems with feedback control is: x k+1 =F (x k , u k ) y k =Z(x k ) u k =h(y 0 , ..., y k ) (7) with state vector x, dynamics function F , measurement outputs y, applied control u, and the subscript denoting the time step. The feedback rule h has access to past and present measurements. A problem in optimal control is to design a control rule, h, that results in a sequence y 0:T . = {y 0 , ..., y T } (or x 0:T directly) that scores well according to some cost function C. Examples include simply reaching a goal condition, C = |y T − y|, or following close to a desired trajectory, y 0:T . Systems with simpler dependence on the input are generally easier to analyze and control (i.e. simpler h performs well), even if the dependence on the state is complicated  (Skelton, 1988) . Control-affine systems are a broad class of dynamical systems which are especially amenable to analysis  (Isidori et al., 1995) . They take the form: F (x k , u k ) = f (x k ) + g(x k )u k (8) where f and g may be nonlinear in state, and are possibly uncertain, meaning unknown. We will seek control-affine form for ease of control and to support future analysis.
 Modified Lagrangian Methods for Constrained Optimization Lagrangian methods are a classic family of approaches to solving constrained optimization problems. We propose an intuitive, previously overlooked form for the multiplier update and derive its beneficial effect on the learning dynamics. We begin by reviewing a prior formulation for the equality-constrained problem. 1 4.1. Review: "Basic Differential Multiplier Method" We follow the development of  Platt & Barr (1988) , who analyzed the dynamics of a continuous-time neural learning system applied to this problem (our result can similarly be derived for iterative gradient methods). They begin with the component-wise differential equations: ẋi = − ∂L(x, λ) ∂x i = − ∂f ∂x i − λ ∂g ∂x i (9) λ = α ∂L(x, λ) ∂λ = αg(x) (10) where we have inserted the scalar constant α as a learning rate on λ. Differentiating (9) and substituting with (10) leads to the second-order dynamics, written in vector format: ẍ + A ẋ + αg(x)∇g = 0 (11) which is a forced oscillator with damping matrix equal to the weighted sum of Hessians: A ij = ∂ 2 f ∂x i ∂x j + λ ∂ 2 g ∂x i ∂x j , or, A = ∇ 2 f + λ∇ 2 g (12) Platt &  Barr (1988)  showed that if A is positive definite, then the system (11) converges to a solution that satisfies the constraint.  Platt & Barr (1988)  also noted that the system (9)-(  10 ) is prone to oscillations as it converges into the feasible region, with frequency and settling time depending on α. We provide complete derivations of the dynamics in (11) and for our upcoming methods in an appendix.
 Proportional-Integral Multiplier Method In (10), λ simply integrates the constraint. To improve the dynamics towards more rapid and stable satisfaction of constraints, we introduce a new term in λ that is proportional to the current constraint value. In the differential equation for λ, this term appears as the time-derivative of the constraint: λ = αg(x) + β ġ(x) = αg(x) + β j ∂g ∂x j ẋj ( 13 ) with strength coefficient, β. Replacing (10) by (  13 ) and combining with (9) yields similar second-order dynamics as (11), except with an additional term in the damping matrix: ẍ + A + β∇g∇ ⊤ g ẋ + αg(x)∇g = 0 (14) The new term is beneficial because it is positive semidefinite-being the outer product of a vector with itself-so it can increase the damping eignevalues, boosting convergence. The results of  (Platt & Barr, 1988 ) hold under  (13, 14) , because the conditions of the solution, namely ẋ = 0 and g(x) = 0, remain unaffected and extend immediately to ġ(x) = 0 (and for the sequel, to g(x) = 0). To our knowledge, this is the first time that a proportional-integral update rule has been considered for a learned Lagrange multiplier. The well-known penalty method  (Hestenes, 1969; Powell, 1969)  augments the Lagrangian with an additional term, c 2 g(x) 2 , which produces a similar effect on the damping matrix, as shown in  (Platt & Barr, 1988) : A penalty = A + c∇g∇ ⊤ g + cg(x)∇ 2 g (15) Our approach appears to provide the same benefit, without the following two complications of the penalty method. First, the penalty term must be implemented in the derivative ẋ, whereas our methods do not modify the Lagrangian nor the derivative in (9). Second, the penalty introduces another instance of the hessian ∇ 2 g in the damping matrix, which might not be positive semi-definite but shares the proportionality factor, c, with the desired term.
 Integral-Derivative Multiplier Method A similar analysis extends to the addition of a term in λ based on the derivative of the constraint value. It appears in λ as the second derivative of the constraint: λ = αg(x) + γg(x) (16) with strength coefficient γ. The resulting dynamics are: ẍ + B −1 A ẋ + αg(x) + γ ẋ⊤ ∇ 2 g ẋ B −1 ∇g = 0 (17) with B = I + γ∇g∇ ⊤ g , and I the identity matrix. The effects of the derivative update method are two-fold. First, since the eigenvalues of the matrix B −1 will be less than 1, both the damping (A) and forcing (∇g) terms are weakened (and rotated, generally). Second, the new forcing term can be interpreted as a drag quadratic in the speed and modulated by the curvature of the constraint along the direction of motion. To illustrate cases, if the curvature of g is positive along the direction of travel, then this term becomes a force for decreasing g. If at the same time g(x) > 0, then the traditional force will also be directed to decrease g, so the two will add. On the other hand, if g curves negatively along the velocity, then the new force promotes increasing g; if g(x) > 0, then the two forces subtract, weakening the acceleration ẍ. By using curvature, the derivative method acts predictively, but may be prone to instability. The proportional-integral-derivative multiplier method is the combination of the previous two developments, which induced independent changes in the dynamics (i.e. insert the damping matrix of (  14 ) into (  17 )). We leave for future work a more rigorous analysis of the effects of the new terms, along with theoretical considerations of the values of coefficients α, β, and γ. In the next section, we carry the intuitions from our analysis to make practical enhancements to Lagrangian-based constrained RL algorithms.
 Feedback Control for Constrained RL We advance the broader consideration of possible multiplier update rules by reinterpreting constrained RL as a dynamical system; the adaptive penalty coefficient is a control input, and the cost threshold is a setpoint which the system should maintain. As the agent learns for rewards, the upward pressure on costs from reward-learning can change, requiring dynamic response. In practical Lagrangian RL, the iterates λ k may deviate from the optimal value, even for lucky initialization λ 0 = λ * , as the policy is only partially optimized at each iteration. Adaptive sequences λ 0 , ..., λ K other than those prescribed by the Lagrangian method may achieve superior cost control for Problem (6). In this section we relate the Lagrangian method to a dynamical system, formalizing how to incorporate generic update rules using feedback. We return to the case of an inequality constrained CMDP to present our main algorithmic contribution-the use of PID control to adapt the penalty coefficient.
 Constrained RL as a Dynamical System We write constrained RL as the first-order dynamical system: θ k+1 =F (θ k , λ k ) y k =J C (π θ k ) λ k =h(y 0 , ..., y k , d) (18) where F is an unknown nonlinear function 2 corresponding to the RL algorithm policy update on the agent's parameter vector, θ. The cost-objective serves as the system measure, y, which is supplied to the feedback control rule, h, along with cost limit, d. From this general starting point, both the RL algorithm, F , and penalty coefficient update rule, h, can be tailored for solving Problem (6). The reward and cost policy gradients of the first-order 3 Lagrangian method, ∇ θ L(θ, λ) = ∇ θ J(π θ ) − λ∇ θ J C (π θ ), can be organized into the form of (18) as: F (θ k , λ k ) = f (θ k ) + g(θ k )λ k (19) f (θ k ) = θ k + η∇ θ J(π θ k ) (20) 2 Known as an "uncertain" nonlinear function in the control literature, meaning we lack an analytical expression for it.  3  We discuss only the first-order case, which provides sufficient clarity for our developments.
 g(θ k ) = −η∇ θ J C (π θ k ) (21) with SGD learning rate η. The role of the controller is to drive inequality constraint violations (J c − d) + to zero in the presence of drift from reward-learning due to f . The Lagrange multiplier update rule for an inequality constraint uses subgradient descent: λ k+1 = (λ k + K I (J C − d)) + (22) with learning rate K I and projection into λ ≥ 0. This update step is clearly an integral control rule, for h.
 Constraint-Controlled RL Our general procedure, constraint-controlled RL, is given in Algorithm 1. It follows the typical minibatch-RL scheme, and sampled estimates of the cost criterion, ĴC are fed back to control the Lagrange multiplier. In contrast to prior work  (Tessler et al., 2018; Paternain et al., 2019)  which uses a single value approximator and treats r + λc as the reward, we use separate value-and cost-value approximators, since λ may change rapidly. When λ is large, the update in (  19 ) can cause excessively large change in parameters, θ, destabilizing learning. To maintain consistent step size, we use a re-scaled objective for the θ-learning loop: θ * (λ) = arg max θ J − λJ C = arg max θ 1 1 + λ (J − λJ C ) This convex combination of objectives yields the policy gradient used in Algorithm 1. Our experiments use this re-scaling, including for traditional Lagrangian baselines.  a ∼ π(•|s; θ), s ′ ∼ T (s, a), 7: r ∼ R(s, a, s ′ ), c ∼ C(s, a, s ′ ) 8: Apply feedback control: Update critics, V φ (s), V C,ψ (s) ⊲ if using 13: ∇ θ L = 1 1+λ ∇ θ Ĵ(π θ ) − λ∇ θ ĴC (π θ ) 14: until converged 15: return π θ 16: end procedure As an aside, we note that it is possible to maintain the control-affine form of (  19 ) with this re-scaling, by reparam-eterizing the control as 0 ≤ u = λ 1+λ ≤ 1 and substituting for (21) with: g(θ k ) = −η∇ θ (J(π θ k ) + J C (π θ k )) (23) This parameterization simply weights the reward and cost gradients in the Lagrangian objective as: ∇ θ L(θ, λ) = (1 − u)∇ θ J(π θ ) − u∇ θ J C (π θ ) (24) It may provide superior performance in some cases, as it will behave differently in relation to the nonlinearity in control which arises from the inequality constraint. We leave experimentation with direct control on u ∈ [0, 1] to future work.
 The PID Lagrangian Method We now specify a new control rule for use in Algorithm 1. To overcome the shortcomings of integral-only control, we follow the developments of the previous section and introduce the next simplest components: proportional and derivative terms. Our PID update rule to replace (  22  return λ
 PID Control Experiments We investigated the performance of our algorithms on Problem (6) in a deep RL setting. In particular, we show the effectiveness of PID control at reducing constraint violations from oscillations and overshoot present in the baseline Lagrangian method. Both maximum performance and robustness to hyperparameter selection are considered. Although many methods exist for tuning PID parameters, we elected to do so manually, demonstrating ease of use.
 Environments: Safety-Gym We use the recent Safety-Gym suite  (Ray et al., 2019) , which consists of robot locomotion tasks built on the MuJoCo simulator  (Todorov et al., 2012) . The robots range in complexity from a simple Point robot to the 12-jointed Doggo, and they move in an open arena floor. Rewards have a small, dense component encouraging movement toward the goal, and a large, sparse component for achieving it. When a goal is achieved, a new goal location is randomly generated, and the episode continues until the time limit at 1,000 steps. Each task has multiple difficulty levels corresponding to density and type of hazards, which induce a cost when contacted by the robot (without necessarily hindering its movement). Hazards are placed randomly at each episode and often lay in the path to the goal. Hence the aims of achieving high rewards and low costs are in opposition. The robot senses the position of hazards and the goal through a coarse, LIDAR-like mode. The output of this sensor, along with internal readings like the joint positions and velocities, comprises the state fed to the agent. Figure  2  displays a scene from the DOGGOGOAL1 environment.
 Algorithm: Constraint-Controlled PPO We implemented Algorithm 1 on top of Proximal Policy Optimization (PPO)  (Schulman et al., 2017)  to make constraintcontrolled PPO (CPPO). CPPO uses an analogous clipped surrogate objective for the cost as for the reward. Our policy is a 2-layer MLP followed by an LSTM with a skip connection. We applied smoothing to proportional and derivative controls to accommodate noisy estimates. The environments' finite horizons allowed use of nondiscounted episodic costs as the constraint and input to the controller. Additional training details can be found in supplementary materials, and our implementation is available at https://github.com/astooke/safe-rlpyt.
 Main Results We compare PID controller performance against the Lagrangian baseline under a wide range of settings. Plots showing the performance of the unconstrained analogue confirm that constraints are not trivially satisfied, and they appear in supplementary material.
 ROBUST SAFETY WITH PI CONTROL We observed cost oscillations or overshoot with slow settling time in a majority of Safety Gym environments when using the Lagrangian method. Figure  3  shows an example where PI-control eliminated this behavior while maintaining good reward performance, in the challenging DOGGOBUTTON1 environment. Individual are plotted for different cost limits. As predicted in  (Platt & Barr, 1988) , we found the severity of cost overshoot and oscillations to depend on the penalty coefficient learning rate, K I . The top left panel of Figure  4  shows example cost curves from DOGGOGOAL2 under I-control, over a wide range of values for K I (we refer to varying K I , assuming K I = 1; the two are interchangeable in our design). With increasing K I , the period and ampli-tude of cost oscillations decrease and eventually disappear. The bottom left of Figure  4 , however, shows that larger K I also brings diminishing returns. We study this effect in the next section. The center and right columns of Figure  4  show the cost and return when using PI-control, with K P = 0.25 and K P = 1, respectively. Proportional control stabilized the cost, with most oscillations reduced to the noise floor for K I > 10 −4 . Yet returns remained relatively high over a wide range, K I < 10 −1 . Similar curves for other Safety Gym environments are included in an appendix.  We examine the trade-off between reward and constraint violation by forming an overall cost figure of merit (FOM). We use the sum of non-discounted constraint violations over the learning iterates, C F OM = k (D(π θ k ) − d) + , D(π θ ) = E τ ∼π T t=0 C(s t , a t , s ′ t ) , and estimate it online from the learning data. Figure  5  compares final returns against this cost FOM for the same set of experiments as in Figure  4 . Each point represents a different setting of K I , averaged over four runs. PI-control expanded the Pareto frontier of this trade-off into a new region of high rewards at relatively low cost which was inaccessible using the Lagrangian method. These results constitute a new state of the art over the benchmarks in  Ray et al. (2019) . We performed similar experiments on several Safety Gym environments in addition to DOGGOGOAL2: POINTGOAL1, the simplest domain with a point-like robot, CARBUTTON1, for slightly more challenging locomotive control, and DOG-GOBUTTON1 for another challenging task (see appendix for learning curves like Figure  4 ). for two strengths of added proportional control, for these environments. PI-control clearly improved the cost FOM (lower is better) for K I < 10 −1 , above which the fast integral control dominated. Hence robustness to the value for K I was significantly improved in all the learning tasks studied. . Learning run cost FOM versus penalty learning rate, KI , from four environments spanning the robots in Safety Gym. Each point is an average over four runs. In all cases, PI-control improves performance (lower is better) over a wide and useful range of KI , easing selection of that hyperparameter.
 CONTROL EFFICIENCY We further investigated why increasing the penalty learning rate, K I , eventually reduces reward performance, as was seen in the robustness study. Figure  7  shows learning curves for three settings: I-and PI-control with the same, moderate K I = 10 −3 , and I-control with high K I = 10 −1 . The high-K I setting achieved responsive cost performance but lower long-term returns, which appears to result from wildly fluctuating control. In contrast, PI-control held relatively steady, despite the noise, allowing the agent to do reward-learning at every iteration. The bottom panel displays individual control iterates, here displayed as u = λ/(1 + λ), over the first 7M environment steps, while the others show smoothed curves over the entire learning run, over 40M steps.  It was further able to slow the approach of the cost curve towards the limit, a desirable behavior for online learning systems requiring safety monitoring. Curves for other environments are available in an appendix.
 Reward-Scale Invariance In the preceding sections, we showed that PID control improves hyperparameter robustness in every constrained RL environment we tested. Here we propose a complementary method to promote robustness both within and across environments. Specifically, it addresses the sensitivity of learning dynamics to the relative numerical scale of reward and cost objectives. Consider two CMDPs that are identical except that in one the rewards are scaled by a constant factor, ρ. The optimal policy parameters, θ * remain unchanged, but clearly λ * must scale by ρ. To attain the same learning dynamics, all controller settings, λ 0 , K I , K P , and K D must therefore be scaled by ρ. This situation might feature naturally within a collection of related learning environments. Additionally, within the course of learning an individual CMDP, the balance between reward and cost magnitudes can change considerably, placing burden on the controller to track the necessary changes in the scale of λ. One way to promote performance of a single choice of controller settings across these cases would be to maintain a fixed meaning for the value of λ in terms of the relative influence of reward versus cost on the parameter update. To this end, we introduce an adjustable scaling factor, β k , in the policy gradient: ∇ θ L = (1 − u k )∇ θ J(π θ k ) − u k β k ∇ θ J C (π θ k ) (25) A conspicuous choice for β k is the ratio of un-scaled policy gradients: β ∇,k = ||∇ θ J(π θ k )|| ||∇ θ J C (π θ k )|| (26) since it balances the total gradient to have equal-magnitude contribution from reward-and cost-objectives at λ = 1 and encourages λ * = 1. Furthermore, β ∇ is easily computed with existing algorithm components. To test this method, we ran experiments on Safety Gym environments with their rewards scaled up or down by a factor of 10. Figure  9  shows a representative cross-section of results from the POINTGOAL1 environment using PI-control. The different curves within each plot correspond to different reward scaling. Without objective-scaling (i.e. β = 1), the dynamics under ρ = 10 are as if controller parameters were instead divided by 10, and likewise for ρ = 0.1. Note the near-logarithmic spacing of λ (λ ρ=10 has not converged to its full value). Using β ∇ , on the other hand, the learning dynamics are nearly identical across two orders of magnitude of reward scale. λ 0 = 1 becomes an obvious choice for initialization, a point where previous theory provides little guidance  (Chow et al., 2019)  (although here we left λ 0 = 0). Experiments in other environments and controller settings yielded similar results and are included in supplementary materials. Other methods, such running normalization of rewards and costs, could achieve similar effects and are worth investigating, but our simple technique is surprisingly effective and is not specific to RL.
 Conclusion Starting from a novel development in classic Lagrangian methods, we introduced a new set of constrained RL solutions which are straightforward to understand and implement, and we have shown them to be effective when paired with deep learning. Several opportunities for further work lay ahead. Analysis of the modified Lagrangian method and constrained RL as a dynamical system may relax theoretical requirements for a slowly-changing multiplier. The mature field of control theory (and practice) provides tools for tuning controller parameters. Lastly, the control-affine form may assist in both analysis (see  Liang-Liang Xie & Lei Guo (2000)  and  Galbraith & Vinter (2003)  for controllability properties for uncertain nonlinear dynamics) and by opening to further control techniques such as feedback linearization. Our contributions improve perhaps the most commonly used constrained RL algorithm, which is a workhorse baseline. We have addressed its primary shortcoming while preserving its simplicity and even making it easier to use-a compelling combination to assist in a wide range of applications. Figure 1 . 1 Figure1. Left: The traditional Lagrangian method exhibits oscillations with 90 • phase shift between the constraint function and the Lagrange multiplier, characteristic of integral control. Right: PID control on the Lagrange multiplier damps oscillations and obeys constraints. Environment: DOGGOBUTTON1, cost limit 200.
 ) is shown in Algorithm 2. The proportional term will hasten the response to constraint violations and dampen oscillations, as derived in Section 4. Unlike the Lagrangian update, derivative control can act in anticipation of violations. It can both prevent cost overshoot and limit the rate of cost increases within the feasible region, useful when monitoring a system for further safety interventions. Our derivative term is projected as (•) + so that it acts against increases in cost but does not impede decreases. Overall, PID control provides a much richer set of controllers while remaining nearly as simple to implement; setting K P = K D = 0 recovers the traditional Lagrangian method. The integral term remains necessary for eliminating steady-state violations at convergence. Our experiments mainly focus on the effects of proportional and derivative control of the Lagrange multiplier in constrained deep RL.Algorithm 2 PID-Controlled Lagrange Multiplier 1: Choose tuning parameters: K P , K I , K D ≥ 0 2: Integral: I ← 0 3: Previous Cost: J C,prev ← 0 4: repeat at each iteration k ← (K P ∆ + K I I + K D ∂) + 10: J C,prev ← J C 11:
 Figure 2 . 2 Figure 2. Rendering from the DOGGOGOAL1 environment from Safety Gym. The red, four-legged robot must walk to the green cylinder while avoiding other objects, and receives coarse egocentric sensor readings of their locations.
 Figure 3 . 3 Figure 3. Oscillations in episodic costs (and returns) from the Lagrangian method, KP = 0, KI = 10 −2 , are damped by proportional control, KP = 1 (ours), at cost limits 50, 100, 150, 200 (curves shaded) in DOGGOBUTTON1.
 Figure 4 . 4 Figure 4. Top row: Constraint-violating oscillations decrease in magnitude and period from increases in the Lagrange multiplier learning rate, KI . At all levels, oscillations are damped by PIcontrol, KP = 0.25, 1. Bottom row: Returns diminish for large KI ; proportional control maintains high returns while reducing constraint violations. Environment: DOGGOGOAL2, cost limit 50.
 Figure 5 . 5 Figure 5. Pareto frontier of return versus cost FOM, which improves (up and to the left) with PI-control, KP = 0.25, 1. Each point is a different setting of KI (see Figure 4).
 Figure6. Learning run cost FOM versus penalty learning rate, KI , from four environments spanning the robots in Safety Gym. Each point is an average over four runs. In all cases, PI-control improves performance (lower is better) over a wide and useful range of KI , easing selection of that hyperparameter.
 Figure 7 . 7 Figure 7. I-and PI-control with moderate KI = 10 −3 and Icontrol with fast KI = 10 −1 (IK I +). Top Returns diminished for fast-KI , but high for PI. Second Cost oscillations mostly damped by PI, removed by fast-KI . Third Control (smoothed) varies more rapidly under fast-KI , is relatively steady for PI. Bottom Control over first 500 RL iterations; fast-KI slams the control to the extremes, causing the diminished returns. Environment: DOG-GOBUTTON1, cost limit 200.
 Figure 8 . 8 Figure8. Derivative control can prevent cost overshoot and slow the rate of cost increase within feasible regions, which the Lagrangian method cannot do. Environment: DOGGOBUTTON1, cost limit 200.
 Figure 9 . 9 Figure 9. Costs, returns, and Lagrange multiplier with rewards scaled by ρ ∈ {0.1, 1, 10}; PI-control with KI = 1e − 3, KP = 0.1. Left column: without objective-weighting, learning dynamics differ dramatically due to required scale of λ. Right column: with objective-weighting, learning dynamics are nearly identical. Environment: POINTGOAL1, cost limit 25.
			 Standard techniques extend our results to inequality constraints, and multiple constraints, as in Platt & Barr (1988) , and notation is simplest for an equality constraint.

Abstract: Lagrangian methods are widely used algorithms for constrained optimization problems, but their learning dynamics exhibit oscillations and overshoot which, when applied to safe reinforcement learning, leads to constraint-violating behavior during agent training. We address this shortcoming by proposing a novel Lagrange multiplier update method that utilizes derivatives of the constraint function. We take a controls perspective, wherein the traditional Lagrange multiplier update behaves as integral control; our terms introduce proportional and derivative control, achieving favorable learning dynamics through damping and predictive measures. We apply our PID Lagrangian methods in deep RL, setting a new state of the art in Safety Gym, a safe RL benchmark. Lastly, we introduce a new method to ease controller tuning by providing invariance to the relative numerical scales of reward and cost. Our extensive experiments demonstrate improved performance and hyperparameter robustness, while our algorithms remain nearly as simple to derive and implement as the traditional Lagrangian approach.