A teacher-student framework to distill future trajectories

1 INTRODUCTION . The ability to learn models of the world has long been argued to be an important ability of intelligent agents . An open and actively researched question is how to learn world models at the right level of abstraction . This paper argues , as others have before , that model-based and model-free methods lie on a spectrum in which advantages and disadvantages of either approach can be traded off against each other , and that there is an optimal compromise for every task . Predicting future observations allows extensive use of all observations from previous experiences during training , and to swiftly transfer to a new reward if the learned model is accurate . However , due to partial observability , stochasticity , irrelevant dynamics and compounding errors in planning , model-based methods tend to be outperformed asymptotically ( Pong et al. , 2018 ; Chua et al. , 2018 ) . On the other end of the spectrum , purely model-free methods use the scalar reward as the only source of learning signal . By avoiding the potentially impossible task of explicitly modeling the environment , model-free methods can often achieve substantially better performance in complex environments ( Vinyals et al. , 2019 ; OpenAI et al. , 2019 ) . However , this comes at the cost of extreme sample inefficiency , as only predicting rewards throws away useful information contained in the sequences of future observations . What is the right way to incorporate information from trajectories that are associated with the inputs ? In this paper we take a step back : Instead of trying to answer this question ourselves by handdesigning what information should be taken into consideration and how , we let a model learn how to make use of the data . Depending on what works well within the setting , the model should learn if and how to learn from the trajectories available at training time . We will adopt a teacher-student setting : a teacher network learns to extract relevant information from the trajectories , and distills it into target activations to guide a student network.1 A sketch of our approach can be found in Figure 1 , next to prototypical computational graphs used to integrate trajectory information in most model-free and model-based methods . Future trajectories can be seen as being a form of privileged information Vapnik and Vashist ( 2009 ) , i.e . data available at training time which provides additional information but is not available at test time . 1Note that the term distillation is often used in the context of “ distilling a large model into a smaller one ” ( Hinton et al. , 2015 ) , but in this context we talk about distilling a trajectory into vectors used as target activations . Contributions The main contribution of this paper is the proposal of a generic method to extract relevant signal from privileged information , specifically trajectories of future observations . We present an instantiation of this approach called Learning to Distill Trajectories ( LDT ) and an empirical analysis of it . x h1 h2 h3 ŷy f1 f2 f3 f4 ( a ) Model-free x x̂1 x̂2 ... x̂n−1 x̂n x∗1 x∗2 x∗n−1 x∗n ŷ1 ŷ2 ŷn−1 ŷn ŷ y1 y2 yn−1 yT Σ f f f f f r r r r ( b ) Vanilla model-based x h1 h2 h3 x̂∗ x∗ŷy f1 f2 f3 f4 faux ( c ) Auxiliary task x x∗ h1 h2 h3 h∗1 h∗2 h∗3 ŷ y f1 f2 f3 f4 T ( d ) Teacher Figure 1 : Comparison of architectures . The data generator is a Markov reward process ( no actions ) with an episode length of n. x denotes the initial observation . y = ∑ i yi is the n-step return ( no bootstrapping ) . x∗ = ( x∗1 , x ∗ 2 , ... , x ∗ n ) is the trajectory of observations ( privileged data ) . Model activations and predictions are displayed boxed . Losses are displayed as red lines . Solid edges denote learned functions . Dotted edges denote fixed functions . 2 RELATED WORK . Efficiently making use of signal from trajectories is an actively researched topic . The technique of bootstrapping in TD-learning ( Sutton , 1988 ) uses future observations to reduce the variance of value function approximations . However , in its basic form , bootstrapping provides learning signal only through a scalar bottleneck , potentially missing out on rich additional sources of learning signal . Another approach to extract additional training signal from observations is the framework of Generalized Value Functions ( Sutton et al. , 2011 ) , which has been argued to be able to bridge the gap between model-free and model-based methods as well . A similar interpretation can be given to the technique of successor representations ( Dayan , 1993 ) . A number of methods have been proposed that try to leverage the strengths of both model-free and model-based methods , among them Racanière et al . ( 2017 ) , who learn generative models of the environment and fuse predicted rollouts with a model-free network path . In a different line of research , Silver et al . ( 2017 ) and Oh et al . ( 2017 ) show that value prediction can be improved by incorporating dynamical structure and planning computation into the function approximators . Guez et al . ( 2019 ) investigate to what extent agents can learn implicit dynamics models which allow them to solve planning tasks effectively , using only model-free methods . Similarly to LDT , those models can learn their own utility-based state abstractions and can even be temporally abstract to some extent . One difference of these approaches to LDT is that they use reward as their only learning signal without making direct use of future observations when training the predictor . The meta-gradient approach presented in this paper can be used more generally for problems in the framework of learning using privileged information ( LUPI , ( Vapnik and Vashist , 2009 ; Lopez-Paz et al. , 2016 ) ) , where privileged information is additional context about the data that is available at training time but not at test time . Hindsight information such as the trajectories in a value-prediction task falls into this category . There are a variety of representation learning approaches which can learn to extract learning signal from trajectories . Jaderberg et al . ( 2016 ) demonstrate that the performance of RL agents can be improved significantly by training the agent on additional prediction and control tasks in addition to the original task . Du et al . ( 2018 ) use gradient similarity as a means to determine whether an auxiliary loss is helpful or detrimental for the downstream task . Oord et al . ( 2018 ) introduce a method based on contrastive learning . They , as well as multiple follow-up studies , show that the representations learned in this way are helpful for downstream tasks in a variety of settings . Buesing et al . ( 2018 ) present ways to learn efficient dynamical models which do not need to predict future observations at inference time . Recently , Schrittwieser et al . ( 2019 ) introduced an RL agent that learns an abstract model of the environment and uses it to achieve strong performance on several challenging tasks . Similarly to our motivation , their model is not required to produce future observations . Meta-learning approaches have recently been shown to be successful as a technique to achieve fast task adaptation ( Finn et al. , 2017 ) , strong unsupervised learning ( Metz et al. , 2019 ) , and to improve RL ( Xu et al. , 2018 ) . Similar to this paper in motivation is the recent work by Guez et al . ( 2020 ) which also investigates how privileged hindsight information can be leveraged for value estimation . The difference to LDT is how the trajectory information is incorporated . Their approach has the advantage of not needing second-order gradients . At the same time , LDT naturally avoids the problem of the label being easily predictable from the hindsight data — the teacher is trained to present it to the student in such a way that it empirically improves the student ’ s performance on held-out data . Veeriah et al . ( 2019 ) use meta-gradients to derive useful auxiliary tasks in the form of generalized value functions . In contrast , we use a teacher network that learns to provide target activations for a student neural network based on privileged information . 3 META-LEARNING A DYNAMICS TEACHER . Here we describe our approach of jointly learning a teacher and a student.2 While our approach applies to the generic setting of learning using privileged information ( Vapnik and Vashist , 2009 ) , here we will focus on the special case of a prediction task with an underlying dynamical system . 3.1 LEARNING TASK . We are considering learning problems in which we have to make a prediction about some property of the future state of a dynamical system , given observations up to the current state . Our method particularly applies to systems in which both the function that relates the current observation to the label as well as the function that predicts the next observation from the current one are hard to learn , making it difficult for both model-free and model-based methods respectively . To make the explanation more concrete , we will use the practical problem of medical decision-making as a running example to which we can relate the definitions we used , similar to a motivating example from Vapnik and Vashist ( 2009 ) : given the history of measurements ( biopsies , blood-pressure , etc . ) on a given patient and the treatment assignment , we want to predict whether the patient will recover or not . The input x ∈ X of our learning task is some observation of the system state st ∈ S before and including time step3 t ∈ Z . In our running example , st can be considered the detailed physical state of the patient , which is not directly observable . The observations x include potentially multi-modal data such as x-ray images , vital sign measurements , oncologist reports , etc . The system is governed by an unknown dynamical law f : S→ S — in our example , the dynamics are physical equations that determine the evolution of all cells in the body . The prediction target y ∈ Y is some function of a future state sT = fT−t ( st ) , separated from t by T − t time steps : y = g ( sT ) . In our running example , a prediction target could be the binary indicator of whether the patient will recover within some time frame . Note that T could vary from one example to the next . In addition to the initial observation , we have access to the trajectory x∗ = ( xτ ) τ=t+1 .. T at training ( but not test ) time . In our running example , the trajectory includes all measurements from the patient after the treatment decision has been made . This information is available in a dataset of past patients ( in hindsight ) , but not in any novel situation . 2Note that unlike in some related work , the teacher in our task is not a copy of the student network , but can have a completely different architecture . 3For simplicity , our dynamical system is time-discrete , but this assumption is not important for what follows . 3.2 SUPERVISION OF INTERNAL ACTIVATIONS . A straightforward approach to solve the learning task which takes into account the trajectory information , would be to train a state-space-model ( SSM ) f̂ , consisting of a dynamical model and a decoder . The SSM is trained to maximize the likelihood of the observed trajectories in the training set , conditioned on the observed initial observation . Ideally , the induced f̂ closely resembles f , such that at test time , we can use it to generate an estimate of the rollout and infer the label from it . A potential drawback of this approach is that learning a full SSM could be more difficult than necessary . There may be many details of the dynamics that are both difficult to model and unimportant for the classification tasks . One example for this is the precise timing of events . As argued by Neitz et al . ( 2018 ) ; Jayaraman et al . ( 2019 ) , there are situations in which it is easy to predict a sequence of events where each event follows a previous one , but hard to predict the exact timing of those events . Moreover , an SSM typically requires rendering observations at training time , which may be difficult to learn and computationally expensive to execute . In the running example from Section 3.1 , it seems challenging and wasteful to predict all future observations in detail , as it would require modeling a complicated distribution over data such as X-ray images or doctor reports written in natural language . Ideally , we would like a model to learn how to extract the relevant information from these data efficiently . We propose to relax the requirement of fitting the dynamics precisely . The teacher can decide to omit properties of the observations that are not needed and omit time steps that can be skipped . It could also change the order of computation and let the student compute independently evolving sub-mechanisms sequentially , even if they evolved in parallel in the actual data generating process . In addition to potentially simplifying the learning problem , this could have the additional benefit of gaining computational efficiency . For example , modeling detailed pixel observations may be computationally wasteful , as argued by Buesing et al . ( 2018 ) and Oord et al . ( 2018 ) .

Summary: This paper proposes a learning framework for predicting the labels of dynamic systems. Unlike existing model-based approaches and model-free approaches, the proposed model takes a middle ground and uses a knowledge distillation-based framework. It uses a teacher model to learn to interpret a trajectory of the dynamic system, and distills target activations for a student model to learn to predict the system label based only on the current observation.