A teacher-student framework to distill future trajectories
Authors: Alexander Neitz
Giambattista Parascandolo
Bernhard Schölkopf
Venue: ICLR 2021 Poster (2021)
Keywords: meta-learning
privileged information

Abstract: By learning to predict trajectories of dynamical systems, model-based methods can make extensive use of all observations from past experience. However, due to partial observability, stochasticity, compounding errors, and irrelevant dynamics, training to predict observations explicitly often results in poor models. Model-free techniques try to side-step the problem by learning to predict values directly. While breaking the explicit dependency on future observations can result in strong performance, this usually comes at the cost of low sample efficiency, as the abundant information about the dynamics contained in future observations goes unused. Here we take a step back from both approaches: Instead of hand-designing how trajectories should be incorporated, a teacher network learns to interpret the trajectories and to provide target activations which guide a student model that can only observe the present. The teacher is trained with meta-gradients to maximize the student's performance on a validation set. We show that our approach performs well on tasks that are difficult for model-free and model-based methods, and we study the role of every component through ablation studies.

Published as a conference paper at ICLR 2021
A TEACHER-STUDENT FRAMEWORK TO DISTILL FUTURE TRAJECTORIES
Alexander Neitz!, Giambattista Parascandolo'.2" Bernhard Schölkopf!.
MPI for Intelligent Systems, Tübingen, 2ETH, Zürich, equal contribution
ABSTRACT
By learning to predict trajectories of dynamical systems, model-based methods can make extensive use of all observations from past experience. However, due to partial observability, stochasticity, compounding errors, and irrelevant dynamics, training to predict observations explicitly often results in poor models. Model-free techniques try to side-step the problem by learning to predict values directly. While breaking the explicit dependency on future observations can result in strong per- formance, this usually comes the cost of low sample efficiency, as the abundant information about the dynamics contained in future observations goes unused. Here we take a step back from both approaches: Instead of hand-designing how trajectories should be incorporated, teacher network learns to extract relevant information from the trajectories and to distill it into target activations which guide student model that can only observe the present. The teacher is trained with meta- gradients to maximize the student's s performance on a validation set. Our approach performs well on tasks that are difficult for model-free and model-based methods, and we study the role of every component through ablation studies.
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, others have before, that model-based and model-free methods lic 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 new reward if the learned model is accurate. However, due partial observability, stochasticity, irrelevant dynamics and compounding errors in planning, model-based methods tend to be outperformed asymptotically (Pong et al.) 20188) 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 OpenAl 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 take a step back: Instead of trying to answer this question ourselves by hand- designing what information should be taken into consideration and how, we let a model learn how to make use of the data. Depending if and how to learn from the trajectories available at training time. We will adopt a teacher-student setting: teacher network learns to extract relevant information from the trajectories, and distills it into target activations to guide a student network] [ next to prototypical computational graphs used to integrate trajectory information most model-free and model- based methods. Future trajectories can be seen as being form of privileged information Vapnik and Vashist (2009). i.c. data available at training time which provides additional information but is not available at test time.
what works well within the setting, the model should learn
A sketch of our approach can be found in Figure [
Note that the term distillation is often used in the context of "distilling large model into smaller one" (Hinton et al. 2015), but in this context talk about distilling a trajectory into vectors used target activations.

---

Published as a conference paper at ICLR 2021
The main contribution of this paper is the proposal of generic method extract relevant signal from privileged information. specifically trajectories future observations. We present an instantiation of this approach called Learning Distill Trajectories (LDT) and empirical analysis of it.
Contributions
MI
fx
f
10
fux
hs
A5
0n
6x
fx
fx
f2
fo
f2
=10011
hg
hg
M
=1101|
9/
11
f:
(a) Model-free
(b) Vanilla model-based
(c) Auxiliary task
(d) Teacher
Figure Comparison of architectures. The data generator is Markov reward process actions) with an episode length of x denotes the initial observation. y >. 3i is the n-step return (no bootstrapping). xx (zi n2 ...... .. 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.
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, its basic form, bootstrapping provides learning signal only through scalar bottleneck, potentially missing
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 bridge the gap between model-frea and model-based methods as well. A similar interpretation can be given to the technique of successor representations (Dayan 1993).
number of methods have been proposed that try to leverage the strengths of both model-free and model-based methods, among them Racanière al. (2077), who learn generative models of the environment and fuse predicted rollouts with a model-free network path. In a different line of research, (iilver et al. (2017) and Oh al. (2017) show that value prediction can be improved by incorporating dynamical structure and planning computation into the function approximators. Guez al. 2019) investigate 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- peeta-gradient approach presented in this paper 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 test time. Hindsight information such as the trajectories in a value- prediccienitiion task falls into this category.
There are a variety of representation learning approaches which can learn to extract learning signal from trajectories. Jaderberg al. (2016) demonstrate that the performance of RL. agents can be
2

---

Published as a conference paper at ICLR 2021
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. (ord 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 variety of settings.
Buesing et al. (2018) present ways to learn efficient dynamical models which do not need to predict future observations 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 RI. (Xu et al. 2018). Similar to this paper in motivation is the recent work by Guez 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 's performance on held-out data. Veeriah 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 student neural network based on privileged information.
3 META- A DYNAMICS TEACHER
Here we describe our approach of jointly learning a teacher and a student 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.
2 While our approach
3.1 LEARNING TASK
We are considering learning problems in which we have 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 mosel- methods respectively.
To make the explanation more concrete, we will use the practical problem of medical Docision-making as rumning 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-pressue etc.) on a given patient and the treatment assignment, we want to predict whether the patient will recover or not.
The input æ E X of our learning task is some observation of the system state s; E S before and including time step" t E E. In our rumning 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 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 ) is some function of a future state 87 = (xx), separated from / by T / time steps: = g(T)). In our running example, 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+)r=t+1.7 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.
2 Note that unlike in some related work, the teacher in our task is not a copy of the student network, hut can have a completely different architecture. For simplicity, our dynamical system time-discrete, but this assumption is not important for what follows.
3

---

network S, parameterized by weights 0, is the network that attempts to predict the quantity of interest y (for instance a cumulative reward or value). Its input In computing ÿ. it produces a sequence of internal activations (hi,..., lww,, one for each of its N hidden layers. Each h2 is a vector whose size is the number of neurons of the corresponding hidden Published as a conference paper at ICLR 2021
y
Ly
SGD
Ln
Validation
y
B2
L
yv
T
Meta-gradient
Visualization of the LDT framework for the special case of n Circled nodes are part the dataset. I denotes the input, x" is the privileged data, y is the label. S the student network with parameters 0, T is the teacher network with parameters Φ.
Figure 2:
3.2 SUPERVISION OF INTERNAL. ACTIVATIONS
straightforward approach to solve the learning task which takes into account the trajectory in- formation, would to train a state-space-mode (SSM) f. consisting 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 clooeey resembles such that at test time, we can use it to generate an estimate of the rollout and infer the label from A potential drawback of this approach that learning 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.] 22:18; Jayaraman et al. (2019), there are situations in which it is easy to predict sequence of events where each event follows 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.
dynamical model and
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 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. 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 al. (0188.
3.3 STUDENT-TEACHES SETUP
We propose student-teacher setup with two neural networks, as shown Figure 2 The student
X, and its output is S(x) = ğ E y
layer. The student's goal is to minimize the generalization loss Ex,y [Cy (5(x), y) for some loss function Cy: y x R.
The teacher network T. parametrized by weights is only used at training time, not at test time. reads the observations the rollout x² = (x+) corresponding to the current training example, and outputs supervision signals (hj hr))
The target activations produced by the teacher's supervision result in another loss for the student, the teaching loss, defined T.,, Ee Lu(ha, hT). La denotes the teaching loss function which, given
4

---

Published as a conference paper at ICLR 2021
pre-activation and supervision signal, produces scalar value. It can be chosen to be any common loss function. Note however that in general, L1, could combine its inputs in an arbitrary way, as long as it is differentiable and produces a scalar. In particular, fı and h2 are not required to have the
same dimensionality. For example, in our specific instantiation described Section 4 h* contains masking weights gate the teaching signal. The total student training loss is Traii aT., Ly- where Ly is the label loss, c.g. the cross-entropy error between predictions and true labels. E R the teaching coefficient, coefficient weighting the losses against each other.
3.4 TRAINING THE TEACHER USING
We train the teacher's weights using the the technique of meta-gradient optimization. This is done as follows: At the beginning of training. we split the dataset into a training and validation set" This split is kept during the entire duration of training. The split ratio is hyperparameter.
The student's weights are updated times using Stochastic Gradient Descent on randomly sampled training batches, resulting updated weights on The student S is then evaluated on validation set, producing validation loss Leat- This validation loss is optimized by the teacher. This validation loss does not contain a term for the internal activation loss, but consists only of the label-loss C ((vott) Weat). The teacher is optimized via the meta-gradient
dLpat Oc dø
OS (Iival; 0x)
d0n
(S(Ival 0n), Duat) dÿ
09
dø
where Fvat and Yeal are the inputs and targets from the validation set.
We omit the summation over individual loss components avoid cluttering the notation. The crucial quantity dda describes how the final student's weights 6, depend on the teachers weights Φ. It can be computed in linear com- putation time and space in the number of steps in the inner optimization loop using automatic differentiation" The meta-gradient is then used for one step of stochastic gradient descent of the teacher's 's weights Φ. The student's weights are reset to what they were at the beginning of the step, since On was only a hypothetical parame- terization used to determine the meta-gradient. Then, the student is actually trained using the newly updated teacher for certain number of steps. In our experiments, every step of meta- training is followed by N steps of training the student's 's weights where N is hyperparameter. Alg. describes the teacher update formally.
Algorithm
Teacher update
Require: Student parameters
Require: Φ Teacher parameters
Require: n Inner-loop lcarning rate
Require: a: Teaching coefficient
Require: n Number inner-loop steps
1: Oo <- 0
b / is unchanged in teacher update
2: for in [1..n} do
3:
x, y next training batch
ú. t- S(z; 0.)
5:
h" + [()=;
6:
La + ha(h,h")
7:
Ly + Ly(ū.v)
8:
Lerain Ly + aLn
9:
O, <- (0: Lirain
10:
Tval Vval <- validation data
11:
Leai + (S 0w), Mval)
12: Update 6 to reduce Lval
b Eq.
4 EXPERIMENTS
We implemented LDT in PyTorch (Paszke et al. 2019) using higher by Grefenstette et al. (019)) Note that all experiments, we distinguish between validation set and a test set. The validation set is used to train the teacher's parameters. Therefore, in order to allow for fair comparison with con-meta-learning baselines, we train baselines with the full training set and for LDT we split this set into a training and a validation portion. The test set is separate from the validation set and is used only passively to track generalization metrics.
Note that the validation set is separate from the test ser, which is an independently sampled datasct used only to evaluate the generalization performance of all methods.
See Baydin al. (2018) for survey of Automatic Differentiation in Machine Learning.
5

---

Published as a conference paper at ICLR 2021
J (O,I)
No teacher
LDT
E N I)
Training set size
11024
512
Inear
semple
Teat Douto
5e-1
1-13
256
P-10
128
1e-1
Input dim.
Inm.
18
0.5 Accuracy 1.0
(8)
(h)
(c)
(d)
Figure 3: (a) Data generation diagram of task A. (b) Test accuracies on task A. For every combination, we report the maximum test accuracy achieved, averaged over 5 random seeds. (c) Data generation diagram of task B. (d) Test losses task B achieved by LDT and the no-teacher baselines with different entropy regularization coefficients B.
4.1 TOY EXAMPLES
Before moving to datasets of dynamical systems, study two toy tasks in order to give better intuition for situations where privileged data can improve learning even though is unavailable at test time, and at the same time examine whether LDT can make use of the privileged data.
Task A This task demonstrates a situation where the privileged information x² predicts the label perfectly and is lower-dimensional than the input x. At the same time, x= is not deterministically predictable from æ. Formally, we construct the distribution over I, x², and y such that the conditional expectation [ [c] for all x, and the conditional entropies H Hy]xz) intended to correspond to real-world setting where we observe training- time privileged data which gives us a oin-eemnnio explanation of the label, but this explanation has been obfuscated by noise and is hence unnecessarily difficult to predict directly (here impossible with deterministic model).
H (y x*) 0. This is
We first sample a D-dimensional input with independent Gaussian components. This vector gets mapped to 2-dimensional vector h using random but fixed linear transformation A E R2x The label is obtained by applying XOR to > 0, The privileged vector x= is constructed by independently sampling another two-dimensional vector which is multiplied with h and concatenated to it. See Figure 3 for diagram and Appendix B. for a detailed description of the dataset.
Using MLP to predict x² from fails because the optimal predictor of x² from always outputs However, in principle LDT can help in this situation: the teacher could learn to invert the stochastic mapping from h to xr" We set up study examine whether LDT automatically discovers a suitable inversion in practice. As student- and teacher-models we use MLPs with one hidden layer each. The teacher gets x² input and produces target activations for the student's hidden layer. To investigate the sample efficiency, we let both the unguided student and LDT learn for grid of different input dimensionalities D and dataset sizes. The achieved test-set accuracies shown in Figure 3, indicating that using a privileged data and teacher makes this learning problem substantially more sample-efficice
Task B In this task (Figure Be), instead of using deterministic labels and stochastic privileged data, we construct deterministic privileged data and noisy labels. A large neural network trained on these labels will tend to fit them exactly. As we discover empirically, this leads to ill-calibrated out-of-sample predictions.
An example for a practical situation where this applies is learning value function in reinforcement learning using Monte Carlo or n-step temporal difference learning: environments and policies are typically stochastic, resulting in noisy empirical value targets. However, the trajectory of observations and actions as privileged data, can in principle explain away part of the influence of chance in the observed value target. We model this situation by providing privileged data transformed view of the logits that were used to sample the target. This transformation is unknown to the learner, and we investigate whether LDT can still use make the student learn well-calibrated mapping.

---

Published as a conference paper at ICLR 2021
We again sample the D components of each input iid. from N( 1). The random linear now transforms I into a d dimeennsio space. The privileged data r E 2d is obtained via another linear transformation B R4, of h. In our experiment D = 128, d, = 4 and dp = 32, and use 1000 training examples. The teacher gets x" as input and only needs supervise the student's output layer. As baseline train the student without teacher or privileged data, but regularize its output predictions by subtracting 3 (ÿ) from the training loss of each example, where H (ÿ) is the entropy of the model's prediction, and is scalar coefficient. As shown in Figure B1 learning the mapping from stochastic labels alone never learns a well well-calibrated map from y. while LDT learns interpret the privileged information at training time, leading to student that generalizes well at test time.
transformation A E R x x
set
4.2 GAME OF LIFE
We performed an additional experiment aimed at evaluating whether LDT can help to extract dynamical information from trajectories. For that reason we created dataset based on the cellular automaton Game of Life by John Conway. The input is random initial state the output is the state of one particular cell after n evolution steps. The privileged data x" consists of the trajectory of n states after the first one. To make the task more difficult, x² is temporally permuted randomly, but consistently across examples. detailed description is provided in Appendix [ As shown in Fig. 4 LDT can learn from the scrambled trajectory and help the student learn with less data than model-free method.
Trajectory (Privileged data x*)
Initial state (x)
1.0
Test accuuaay
Model free
0.9
Perfact teacher teacher
0.7
Label (y)
128
256
512
1024 2048
Training
Figure 4: A datapoint in the Game-of-L Life task (left) and the results of the experiment (right).
4.3 MUJoCo
In order to evaluate whether LDT can improve learning efficiency in continuous control tasks, set up a prediction task using the MuJoCo simulator (Todorov et al.) 2012). As an objective we choose learning n-step reward model: such model has predict the sum of rewards along trajectory of n steps, given only access to the current state and n (open-loop) actions. We fix n = 16 in all experiments.
We collect a dataset of size 4000 for each of the environments Swimmer-v2, Walker2d-v2, Hopper-v2, and HalfCheetah-v2. For each training example, the input x (St; 0t:t+n) contains both initial state and an action sequence. The initial state 9; is obtained by executing a random policy for short number of time steps after resetting the simulator. The action sequence oc:t+n consists of random actions. The label is the cumulative return of exectuting actions 04:4+ in state See Appendix B.2 for more details on the dataset.
Models
We evaluate the following methods:
Model-free (MF): This baseline resembles the architecture shown in Fig. la We use five-layer fully connected MLP with ReLU activations and 128 neurons per layer. This baseline trained with x as input and y output.
Auxiliary task (Aux): As shown in Fig. Ic this baseline augments the model-free method with an auxiliary task-head, which is trained to fit the full trajectory x* in order to shape the model's internal representation.
LDT: Using the same MLP architecture in the model-free baseline for the student, LDT additionally uses a teacher, which uses a network with a ID-convolutional torso and an MLP-head embed the trajectory xt and provide training signals for the student activations. The teaching
7

---

Published as a conference paper at ICLR 2021
Swimmer-v2
Walker2d-v2
Hopper-v2
HalrCheetah-v2
0.03
0.04
0.006
Test MSE
0.06
MF
0.02
0.03
0.004
Aux
0.05
LDT
0.002
0.01
10
50
100
10
50
100
10
50
100
10
50
Epoch
Epoch
Epoch
Epoch
Figure 5: Test losses for the MuloCo reward prediction task. Evolution of mean squared error between predicted and true normalized n-step reward held-out test set. We ran each configuration with 8 different random seeds and show the aggregated eurves.
loss is a gated mean squared eπror La x Σx (mm)) hg)?, where a is the logistic sigmoid function, m and h² are the outputs by the teacher, and are the student's internal pre-activations.
For all networks follow Schrittwieser al. (2019) in how turn the regression task of predicting rewards into a classification task by binning the reward space (see Appendix for details). Hyperpa- rameters for each method optimized independently (see Appendix for ranges) for each method and task. We select the configuration with the lowest rean-squared-erre on test data and re-run it 8 times with different random seeds.
Results In F we report the mean squared error between predicted and true cumulative reward. The student trained with LDT achieved lower MSE than both the MF and Aux baselines in all tasks. Moreover, found the generalization gap to be significantly smaller for LDT (Figure 66. This can be explained as follows: If the student is overfitting to the training set, its performance will degrade on the validation set. Since this loss only affects the teacher, the teacher can provide teaching targets h" that steer the student away when it starts to severely overfit to the training data.
We acknowledge that there could be strong baselines from the literature that we have not considered. Many of these approaches have complex pipelines of operations (e.g. Chua et al. (2188), or use domain-specific knowledge extract good representations (e.g. CPC (Oord al. 2018)). ). As the
main point of this paper is to investigate and understand new framework to incorporate trajectory information, we decided keep our evaluation setting simple and consistent, e.g. by using the same student architecture for all models.
4.4 A ABLATIONS
Walker2d-v2
We perform several ablation studies, in order to test the role of every component in set-up. We describe every experiment set-up and show the results for the
3x
MF
Test MSE
LDT
environment Walker2d in Figure 7 Results for the
2
other environments are consistent with these, and are shown in Appendix [.2.4
PT
RT-8 RT-16 AR
FT
Ablation type
In the following list, we describe the different ablation studies detail.
Figure 7: Ablation study results
Swimmer-v2
Walker2d-42
Hopper-v2
HalfCheetah--2
Generalization gap
0.003
0.02
0.002
0.05
MF
0.02
0.01
Aux
0.001
LDT
0.00
0.00
0.000
0.00
10
50
100
10
50
100
100
10
Epoch
Epoch
Epoch
Epoch
Figure 6: Generalization gaps (Test-set MSE minus Training-sor set MSE) for the different approaches and domains.)
8

---

Published as a conference paper at ICLR 2021
(i) Fixed untrained teacher (FT): The teacher still provides target activations for the student, but keep the teacher's 's weights fixed to the initial random weights. Using no meta vetta-gradient updates, the student essentially fails to learn, validating the need to train the teacher.
(ii) Time-permuted x² (PT): we independently and randomly permute the time order of x for each example. Interestingly, the performance only degrades slightly. Two possible interpretations of this result the following: either the teacher architecture does not make sufficient use of the temporal structure, or this structure is not so relevant for the task. Since we did not heavily tune our teacher network, we tend to lead towards the first explanation.
(iii) Random x² (RT-n): each entry of the last frames of all trajectories x" replaced by noise sampled i..d. from (000 1). When the trajectories are completely irrelevant to the task, the teacher can in principle learn ignore them. However, since LDT trains with smaller effective training set (because portion is split off for validation), we expect a slightly weaker performance. Indeed, the results show that LDT with fully randomized trajectories (RT-16) tends be comparable to, but slightly worse than the model-free network.
(iv) Same training and validation data (AR): instead fixing the split of the training data into disjoint subsets for the student and for the teacher, we resample the split after every step of meta-training. Results show that the performance is even slightly worse than the model-free baseline. This is consistent with expectations, as the teacher's loss can be minimized directly by the student in training.
CONCLUSION
In this paper, instead of proposing a new hand-designed strategy to incorporate information from previous trajectories, we proposed to let model learn if and how to use them. A teacher network learns to extract relevant information from a trajectory, and distills targets activations for student network that only sees the current observation. The teacher is rewarded for maximizing the student performance on validation data, but can only achieve this indirectly by supervising the student while it trains the training data. The aim of this method is preserve advantages from both model- as well as model-free methods: using the rich amount of information from observations as training signal instead just reward signal, but not capped in performance due to model bias. One advantage that is not preserved from model-based approaches is the straightforward possibility to change tasks by adapting the reward function only. As validated empirically by the experiments and ablations presented in Section 4 this framework allows the model to choose where to sit in the wide spectrum between model-based and model-free methods, adapting to the specifics of the task at hand. An obvious drawback of LDT is like many algorithms that learn how to learn computing meta- gradients increases the time and space complexity at training time by a factor that is linear in the number of inner steps. However, the computational cost of the student at test time is exactly the same as for an equivalent student that did not make use of the privileged information at training time, as the teacher does not play any role and can be discarded.
that
We believe that the general framework of teacher-student trained with meta- gradients to incorporate privileged information can be a fruitful direction for future work, beyond learning from trajectories. As the main limitation is the linear increase in time and space complexity at training time, increases in computing power should allow for more and more complex teachers and students to be trained on large tasks.
ACKNOWLEDGMENTS
We wish to thank Lars Buesing and Arthur Guez for the fruitful discussions on the topic of this project. We thank the International Max Planck Research School for Intelligent Systems for supporting Alexander Neitz, and the Max Planck ETH Center for Learning Systems for supporting Giambattista Parascandolo.

---

Published as a conference paper at ICLR 2021
REFERENCES
A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind. Automatic differentiation in machine learning: a survey. Journal of Machine Learning Research, 18(153):1 2018. URL
http:// jmlr org/ wapers/vl8.... 1 468.html
Bergstra, R. Bardenet, Y. Bengio, and B. Kégl. Algorithms for hyper-parameter optimization. In Pro- ceedings of the 24th International Conference Neural Information Processing Systems, NIPS' 11, page 2546-2554, Red Hook, NY, USA, 2011. Curran Associates Inc. ISBN 9781618395993.
G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. Openai gym, 2016.
L. Buesing, T. Weber, S. Racaniere, S. Eslami, D. Rezende, D. P. Reichert, F. Viola, F. Besse, K. Gregor, D. Hassabis, et al. Learning and querying fast generative models for reinforcement learning. arXiv preprint 1rXiv:180..03006, 2018.
K. Chua, R. Calandra, R. McAllister, and S. Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Biaa-Bianchi,aoo and R. Garnett, editors, Advances in Neural Information Processing Systems 31, Curran Associates, Inc., 2018.
pages 4754-4765.
P. Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Comput,, 5(4): July 1993. ISSN 0899-7667. doi: 10.1 (..1166/ncco.1999 1993.5.4.99. URL
202///cc..... 1933. 5. 4. 613
Y. Du, W. M. Czarnecki, S. M. Jayakumar, R. Pascanu, and B. Lakshminarayanan. Adapting auxiliary losses using gradient similarity. arXiv preprini arXiv: 1812.02224, 2018.
C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pages 1126 1135, International Convention Centre, Sydney, Australia, 06-11 11 Aug 2017. PMLR. URL http: //proceedings.mix prwsspesspeee.ee
htmml
E. Grefenstette, B. Amos, D. Yarats, P. M. Htut, A. Molchanov, F. Meier, D. Kiela, K. Cho, and S. Chintala. Generalized inner loop meta learnng.. arXiv preprint 190.01191. 2019.
A. Guez, M. Mirza, K. Gregor, R. Kabra, S. Racaniere, T. Weber, D. Raposo, A. Santoro, L. Orseau, T. Eccles, G. Wayne, D. Silver, and Lillicrap. An investigation of model-free planning. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2464 2473, Long Beach, California, USA, 09-15 Jun 2019. PMLR. URL http:h::ppoooeednnggge. mmr
html
A. Guez, F. Viola, T. Weber, L. Buesing, S. Kapturowski, D. Precup. D. Silver, and N. Heess. Value-driven hindsight modelling. arXiv preprint arXiv:2002.08329, 2020.
G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in neural network. In NIPS Deep Learning and Representation Learning Workshop, 2015. URL httt:://aaxii// org abs/
1503.021 1503.02531
02531
loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference International Conference on Machine Learning Volume 37, ICML' 15, page 448-456. JMLR.org, 2015.
M. Jaderberg, V. Mnih, W. M. Czarnecki, T. Schaul, J. Z. Leibo, D. Silver, and K. Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.0552 2016
D. Jayaraman, F. Ebert, A. Efros, and S. Levine. Time-agnoster prediction: Predicting predictable video frames. In International Conference on Learning Representations, 2019. URL https 7/openrevien net. wwt.://ooooommmoosyao.aaaoa..
10

---

Published as a conference paper at ICLR 2021
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In Y. Bengio and Y. LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA. May 7-9, 2015, Conference Track Proceedings, 2015. URL httt:://axxiv/ 1412. 6980]
abs/
D. Lopez-Paz, B. Schölkopf, L. Bottou, and V. Vapnik. Unifying distillation and privileged informa- tion. In International Conference on Learning Representations (ICLR), Nov. 2016.
L. Metz, N. Maheswaranathan, B. Cheung, and J. Sohl-Dickstein. Learning unsupervised learn- ing rules. In International Conference on Learning Representations, 2019. URL https
//openreview www.:nee/eeee ienerrrrieee.ntttfoooulee-----e htte://ooommmm--ddddad.a-d
A. Neitz, G. Parascandolo, S. Bauer, and B. Schölkopf. Adaptive skip intervals: Temporal abstraction for recurrent dynamical models. In Advances in Neural Information Processing Systems, pages 9816-9826, 2018.
J. Oh, S. Singh, and H. Lee. Value prediction network. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30, pages 18-6128. Curran Associates, Inc., 2017. http::
nips. https.cc/aapeee http://ppppert111eeeeaeeeeeeeeee
92 www.waaweeeeeeeeeeeo wtwreediiiionnnee
pdf.
A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding, arXiv preprint arXiv: 1807.0307, 2018.
OpenAI, Akkaya, M. Andrychowicz, M. Chociej, M. Litwin, McGrew, Petron, A. Paino, M. Plappert. G. Powell, R. Ribas, Schneider, N. Tezak, J. Tworek, P. Welinder, L. Weng, Q. Yuan, W. Zaremba, and L. Zhang. Solving rubik's cube with a robot hand. CoRR, abs/1910.07113, 2019.
URL httptttpp///art
org/ abs 1110 07113
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf. E. Yang, Z. Detito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages Curran Associates, Inc., 2019.
V. Pong, S. Gu, M. Dalal, and S. Levine. Temporal difference models: Model-free deep RL for model- based control. In International Conference on Learning Representations, 2018. URL
net
S. Racanière, T. Weber, D. Reichert, L. Buesing, A. Guez, D. Jimenez Rezende, A. Puig- domènech Badia, O. Vinyals, N. Heess, Y. Li, R. Pascanu, Battaglia, D. Hassabis, D. Silver, and D. Wierstra. imaginattion-augmentee agents for deep reinforcement learn- ing. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30, pages nips. cc/paperrc/ppppr/
5690-5701. Curran Associates, Inc., 2017. URL for-s--oo---e -op----ooee------e-eeeeeeeeeeeee
7152- 1112-imaginntiin augment agents- pdf
T. Salimans and D. P. Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. In D. D. Lee. M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems 29, pages 901-909. Curran Associates, Inc., 2016.
Schritwieeser, Antonoglou, T. Hubert, K. Simonyan, L. Sifre, Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. arXiv preprint arXiv: Ttt:::991.00066. 108265, 2019.
D. Silver, H. van Hasselt, M. Hessel, Schaul, A. Guez, T. Harley, G. Dulac-Amold D. Re- ichert, N. Rabinowitz, A. Barreto, and T. Degris. The predictron: End-to-end learning and planning. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Confer- ence Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3191-3199, International Convention Centre, Sydney, Australia, 06-11 Aug 2017. PMLR. URL http:://roocet ngs press v00 si ver 17a.html
11

---

Published as a conference paper at ICLR 2021
F. P. Such, A. Rawal, J. Lehman, K. O. Stanley, and J. Clune. Generative teaching networks: Accelerating neural architecture search by learning to generate synthetic training data. arXiv preprint
2019.
R. S. Sutton. Learning predict by the methods of temporal differences, Mach. Learn., 3(1):9-44, Aug. 1988. ISSN 0885-6125, doi: 10.
URL htttps://oi..rrg
I023/A:1022633511
R. Sutton, J. Modayil, M. Delp. T. Degris, P. M. Pilarski, A. White, and D. Precup. Horde: scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems Volume
2, AAMAS 11, page 761-768, Richland, SC, 2011. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 0982657161.
E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In IROS, pages 5026-5033, IEEE, 2012. ISBN 778-1-4673--1737-5
V. Vapnik and A. Vashist. new learning paradigm: Learning using privileged information. Neural networks, (25-6)):544-555 557,, 2009.
V. Veeriah, M. Hessel, Z. Xu, J. Rajendran, R. L. Lewis, J. Oh, H. P. van Has- selt, D. Silver, and S. Singh. Discovery of useful questions as auxiliary tasks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché Buc, E. Fox, and R. Gar- nett, editors, Advances in Neural Information Processing Systems 32, pages 9310- 9321. Curran Associates, Inc., 2019. URL http: attt:://papers nips. cc/paper/ 9129-aa df -ustosiutis as aux auxi wuxiliarr-tasks...
questions
O. Vinyals, I. Babuschkin, W. M. Czarnecki, M. Mathieu, A. Dudzik, J. Chung, D. H. Choi, R. Powell, T. Ewalds, Georgiev, al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(77822:350-354, 2019.
Z. Xu, H. P. van Hasselt, and D. Silver. Meta-gradient reinforcement learning. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gar- nett, editors, Advances in Neural Information Processing Systems 31, pages 2396- 2407. Curran Associates. Inc., 2018. URL /ttppr//ppprr....e www.ssniiis... cc/paper/ 7507 mett-mmtaa-- grwdieaadan..ee reinforcement learning.pdf
12

---

Published as a conference paper at ICLR 2021
A IMPLEMENTATION DETAILS
We implemented Meta-LUPI using Py PyTorch (Paszke et al. 2019) and used the library higher (Grefen
stette al. 2019) compute meta-gradients.
Our implementation of LDT contains number of hyperparameters and settings which we describe here.
'he optimizer used the student's weights 0g in the inner loop of determining the meta-gradient (includes inner learning rate and inner momentum)
Inner-lopp optimizer:
Student optimizer: The optimizer used to actually update the student's weights 0. (Includes student learning rate and student momentum)
Meta-optimizer: The optimizer used optimize the teacher's weights Φb (Includes meta learning rate and meta momentum)
Validation split: The fraction of the training set that used to compute the meta-loss.
Teacher architecture: The architecture of the teacher network.
B EXPERIMENT DETAILS
B.1
TOY DATASETS
Task A
Construction of a training example:
Sample IR. ((0,1
Sample A E R with Aij f...d. Uniform( vD )
(default for torch, nn. Linear)
Sample E (-1,1)2 ftd 2 Bernoull((0.5) -1
h Aæ
1// > 0] + 1[/ > 0] where @ denotes logical XOR.
x² [81, 82; s1h,, s2h2|
Observed at training time: (x, ",,x",y)
Observed at test time: only z.
We perform an experiment over grid of values for and n, where is the number of training examples in the dataset. The range of values for D is the set (64, 128, 256, 512). The range of values for are exponentially-spenn
integers between 128 and 1024. The student is multi-layer- perceptron (MLP) with 32 and 128 neurons in the hidden layers, respectively. The architecture mirrors the data- gotta-genraaiig process with a bottleneck after the first linear transformation in particular, there no activation function after the first hidden layer. We up baseline without teacher, letting the student learn to predict y from x. For the teacher in LDT, we use an MLP with [256, 256] neurons in the hidden layers, ReLU activations, and output size of 32, matching the size of the student's first hidden layer. The teaching loss is the mean squared error between the teacher's output and the student's first hidden activation. Additionally, the student is always trained to minimize binary cross-entropy between its prediction and the training label. We use Adam with learning rate of 10-3 as the optimizer for the student in both methods. For LDT, we use validation split of 0.5.
Task B Construction of training example:
Sample E IRP x ~ 1(,
Sample A E Rddx X with Aig i.idd Uniform( 1
VD )
Sample B E Xds with B, sid. Uniform( Vd. Ton
h Axr
13

---

Published as a conference paper at ICLR 2021
Sample (0..d, 1}: y (ategorical(Sotmmxx(()
r
Bh
Observed at training time: (a, ,,,x"))))
Observed time: only z.
For this experiment, we set D = 128 d, 4, dp = 32, and use 1000 training examples.
The student an MLP with input dimension D, two layers 256 hidden units each, ReLU activations, and output dimension d, The teacher is an MLP with input dimension dp, two layers of 256 hidden units each, ReLU activations, and output dimension da. As teaching loss, we use the KL-divergence between the softmax of the teacher's output and the softmax of the student's 's output activations.
Other hyperparameters used in this experiment are summarized in Table
Table 1: Hyperparameterss for LDT in toy task
Hyperparameter
Value
Inner optimizer
Adam
Inner learning rate
1e-3
Batch size
32
Meta-optimizer
Adam
Meta-learning rate
le-3
Teaching coefficient a
10
Inner loop optimization steps n
Validation split
0.5
B.2
MuJoCo
We now describe additional details for the MuloCo reward-prediction experiments.
B.2.1 DATASET DETAILS
The datasets for the MuloCo reward prediction task are generated as follows.
We use the MuJoCo environments implemented in OpenAl gym Brockman et al. (2016), Each training set consists of 4000 examples generated as follows.
Reset the environment to an initial state drawn from the initial-tatee-disiibbuio
Sample an integer n uniformly from the range (10..30}
Perform n steps following random policy x.
Record the current state s.
Sample 16 step open-loop action sequence 01.16 from
Set the example's 's input (s, a1:16)
Execute 41:16 starting from the current state s
Record the trajectory states 81:16 and rewards T1:16 as privileged data x*
Record the of rewards label y Σ.
As random policy choose the policy that ignores the state and at each step and for each action dimension, samples one the numbers 0- }} with equal probability, independently of each other.
We did not investigate the effects of using a different policy æ to generate the dataset.
The test sets consist of 10k examples distributed identically to their respective training sets.
Similarly to Schrittwieser et al. (2019) we turn the regression task predicting rewards into a classification task by binning the reward space. We first obtain transformation R - [0, 1] in
14

---

Published as a conference paper at ICLR 2021
such way that it transforms training rewards to a uniform distribution in [0, 1], interpolating linearly in-between values from the training set and clipping values outside the seining-reward-aae lie between 0 and 1. We apply this transformation to both training and test labels and afterwards distribute the resulting values into 32 equally spaced bins between O and to obtain categorical values. This leads to even label distribution on the training set, and roughly-even distribution the is then the cross-entropy between predicted label probabilities and
to
test set. The label-loss-functte the one-hot distribution of the true label. The normalized mean-squared-rroor reported in the curves is obtained by first obtaining the expected value of the prediction by weighting the output-bucket bucket with the predicted probability and then determining the squared difference to the true bucket value between 0 and 1.
Before training, we standardize all inputs along the dimensions individually, using empirical means and standard deviations found in the training set. Each state along the trajectory is standardized with the same normalization parameters.
B.2.2 CHOOSING HYPERPARAMETERS
In order to determine good general ranges hyperparameten we first performed manual hyperpa- rameter investigations and used Tree-structured Parzen Estimator (Bergstra et al. (011 search to find good ranges for the hyperparameters. Hyperparameter searches are performed using different random seeds from the final evaluation in order to reduce overfitting due to the hyperparameter- optimization. The curves for LDT shown in Figure Suuure]Jarr are the result re-running the best hyperpa- rameter configurations eight times with different random seeds. The curves for the baselines obtained by performing a grid search on the hyperparameters shown in Tables and selecting the configuration that yielded the lowest test-MSE any point in training and re-running eight times with different random seeds.
We fix as architecture for the prediction model (the student in the LDT framework) a five-layer fully connected MLP with ReLU activations and 128 neurons per layer.
LDT parameters Tables 5 and show hyperparameters we determined to work well for the MuJoCo reward prediction task.
As optimizers for the teacher and student weights, we use Adam (Kingma and Ba 2015). For the inner-loop student-optimizzr we use SGD with momentum.
As the teacher network we use a ID-convolutional Neural Network followed by a fully connected network follows: The 16 states of the trajectory x" (including rewards at every step) are fed into Recovv ID-convolutional layers, each followed by a ReLU activation. The first nccaav 1 convolutional layers have ei output-channels, the last one has 02 output-channels. The output of the last convolutional layer is flattened and fed into a fully connected downstream model with one layer of c3 neurons.
In our experiments we found that it did not help to feed the actions along the trajectory to the teacher (additionally to the student who always gets to see the actions).
The output of the teacher is a vector that contains two entries for every neuron in the student network. One of these values, h is the preattpre-activation, and my is gating parameter weight that particular neuron's 's loss. We found it helpful to scale and shift the input to the gating-sigmoid by two constant scalars that are used for the entire network: the loss for given pre- aretiiivaaionn is computed as a m2 student's network, hi is the target activation given by the teacher, and my the gating signal output by the teacher.
2 1) (hg h )a, where a is the logistic sigmoid function, h& is the pre- rreactivvtionn in the
We use batch-normalization (Toffe and Szegedy, 2015) in the teacher network and weight- normalization (Salimans and Kingma| 2016) in the student network, following the findings of Such et al. (2019).
As precaution against exploding meta- veta-gradients, we clip each component of the meta-gradient to the range [e,,, 1e8), but did not investigate whether this was necessary.
"B hyperparameters we mean those parameters which are fixed over the course of training run. They do consist of the teacher's meto-parameters.
15

---

Published as a conference paper at ICLR 2021
Table 2: Hyperparameter ranges for the method 'model-free" (MF)
Hyperparameter
Range
Learning rate Weight decay
{le- 3,e-3, 1e- 1e-2) 2}
{0, le 10,1e-5,11 1e-4] {8, 32)
Batch size
Table 3: Hyperparamettrr ranges for the method 'auxiliary' (AUX)
Hyperparameter
Range
Learning rate
{1c- 1e 1e-2] {0, 10lle-5. 1e-4]
Weight decay
Auxiliary task weight
{1,
Table 4: Hyperparameter ranges for our proposed method (LDT)
Hyperparameter
Range
31 of Meta-optimizer (Adam)
(0.0,
Learning rate of Inner-loop optimizer (SGD) Teaching coefficient (log10) Validation split
{le 3, 3, 1c-2
{2.0, 2.5, 3.0
{0.3, 0.5, 0.7]
Table 5: Best hyperparameters found in grid search for LDT
Hyperparameter
Swimmer x2
Walker2d x2
Hlopper-22
HaliChestah
31 Meta-optimizer (Adam)
0.9
0.9
0.0
0.0
Learning rate of Inner-loop optimizer (SGD)
le-2
le-2
le-2
le-2
Teaching coefficient 1l0g Validation split
2.5
3.0
2.5
2.5
0.3
0.5
0.3
0.7
Table 6: Other hyperparameters for LDT (all environments)
Hyperparameter
Value (all envs)
Learning rate of Meta-optimizer (Adam)
5e
Learning rate of Student-optimizer (Adam)
le-3 3
Batch size
Momentum of inner-loop optimizer (SGD)
0.75
Weight decay (L2) coefficient in inner-loop
1e-8
Number of inner-loop steps (Ninner)
96
Teacher C1
96
Teacher 22
256
Teacher C3
768
Teacher Reonv
Number student training steps per meta-step
24
Performances When using the hyperparameters described above, obtain the minimum test. MSEs in the eight evaluation runs shown in Table 7
B.2.3 MODEL-BASED BASELINES
naive implementation of model-based could not fit the data. We trained the model within the same teacher-student framework, with a fixed dummy teacher that supervises the output of the student with the next observation (essentially mimicking basel setting). We 'stacked' the student such that it takes input its output observation the previous timestep, and by feeding it only
16

---

Published as a conference paper at ICLR 2021
Table 7: Minimum squared errors all environments
no-teacher
auxiliary
LDT
HalfCheetah- v2
0.0485
0.0477
0.0427
Hopper-v2
0.00270
0.00224
0.00132
Swimmer-v2
0.0169
0.0146
0.00964
Walker2d-v2
0.0238
0.0211
0.0164
Swimmer-v2
Walker2d-v2
10
MF
3x10
MF
Test MSE
2x 10
LDT
Test MSE
LDT
2 10
10
PT
RT-8 RT-16 AR
FT
PT
RT-8 RT-16 AR
FT
Ablation type
Ablation type
Hopper-v2
HaifCheotah--2
4 10
8 10
MF
MF
Test MSE
x 10
LDT
Test MSE
E1 7 10
LDT
6 10
x
5
PT
RT-8 RT-16 FT
PT
RT-8 RT-16
FT
Ablation
Ablation type
Figure 8: All results of ablation studies.
the first observation of the sequence, we supervise all of its intermediate targets. However, without introducing the extra inductive bias of training the student on every step independently of the others (instead of a single trajectory), the student's ' output would tend to diverge as it got deeper into the number of steps.
Training the student pairs of consecutive transitions independently of the whole trajectory, makes the model work much better. However, making a fair comparison to the model-free, Aux, and LDT students is difficult, since the model-based student effectively uses 16 times more computation. Note that the main objective this paper is to compare to what extent the abstract models implicitly learned by the same student architectures but with different techniques, can learn to incorporate the trajectory information. On this basis, we exclude the model-based basclines from our comparisons.
B.2.4 ABLATION STUDIES
In Figure g we present the results of the ablation studies described in Section 4.4
C GAME OF LIFE EXPERIMENT
Here provide details about the Game of Life experiment mentioned in Section 4.2
Dataset As the underlying dynamical system, we use the cellular automaton Game of Life by John Conway. The rules of this cellular automaton are
1. Any live cell with fewer than two live neighbours dies, as if by underpopulation.
"Copied from ht httpm/////a/////./.ooooo/o/////s
17

---

Published as a conference paper at ICLR 2021
2. Any live cell with two three live neighbours lives on to the next generation.
3. Any live cell with more than three live neighbours dies, if by overpopulation.
4. Any dead cell with exactly three live neighbours becomes a live cell, as if by reproduction.
Note that there is no linear decision boundary to determine the next cell state if the input space is the cell'ss neighborhood. However, with at least one hidden layer, it is possible to implement the rules. We therefore dynamics unrolled over n steps.
convolutional neural network with 2n layers to represent the system's
We generate binary classification task from the system follows.
A board of size 17 17 is initialized by setting each cell to "alive" with probability of
Ps 0.4, independently of other cells. Alive cells are represented with a 1, all other cells are represented with O. This first board state is the input {0, 17 for the task.
The Game-of-Life rules are applied three times sequence, yielding three consecutive states
The classification target the middle cell the last state, i.c. a binary label
All intermediate and final states in the trajectory are considered privileged information.
The initial alive- hlive-probability was chosen such that there typical rollouts are diverse. If the initial configuration much sparser much denser, then the population quickly dies off.
Network architecture early, unstructured experiments it seemed as though layers per step make it casier to train CNN using internal supervision. Therefore, we fix three convolutional layers per step for all experiments.
The student's ' internal activations z are convolutional feature maps. We use Z to denote the l'th feature map of the k'th layer. Each z has the same dimensions the map of the cellular automaton (17 x 17 in our experiments).
Implementation of the teacher We choose simple parameterization of the teacher. The teacher's weights are a three-dimensional tensor & E Txxxx where T is the fixed number of steps in the Game-of-Life trajectory, N is the number of convolutional hidden layers in the student network, and F is the number feature maps per hidden layer.
The internal activation targets linear combinations of the cell states in the trajectory (I1, PT):
f
a
(2)
where Aww = softmax( i.e. the mixture weights are the teacher's weights softmaxed- thhoouug time. More explicitly,
exp(óri)
aski Σ, exp((rr1)
(3)
We use Adam to update the inner weights and vanilla SGD+momentum to update the meta-parameters, as we found these two choices generally perform best.
Baselines
Fixed oracle teacher: & is initialized to the values such that it provides the frames in the correct sequence to every third convolutional layer. It is held fixed over the course of training.
Fixed random teacher: & is initialized randomly and held fixed over the course of training.
No teacher: Classification task without privileged data
18

---

Published as a conference paper at ICLR 2021
Experimental protocol The questions that this experiment should help answer is: Does meta- gradient training help with sample efficiency in the Game-of-Life task compared to using (a) no teacher, (b) fixed random teacher?
Training sizes: (128, 256, 512, 1024, 2048, 4096, 8192). but in order to be resource- efficient, try the approaches from fewer examples to more examples until they "max out".
Methods: meta-learned teacher (meta), fixed random teacher (fixed-teacher), perfect teacher (perfect-taache)) no teacher (no-teacher)
In meta-learning, the sppnn-validaioon-slitt happens within the training set (no additional validation data is provided).
We use the same hyperparameter ranges for all experiments ones that seem reasonable based on initial experimentation.
Train for 200 epochs.
Hyperparameter
Distribution
Min
Max
Methods
learning_rate
LogUniform
10
10 2
all
adam_momentum
LogUniform
10
10
all
weight_decay
LogUniform
10
10
all
batch_siie
DiscreteUnifond
32
128
all
weighh__init_mutiiplier
LogUniform
0.1
2.0
all
n_filters
DiscreteUniform
32
internal_coef
LogUniform
10a
10*
meta, fixed-teacher
teacher_wcight_scalle
Log LogUnifora Uniform
10
10
fixed-teacher
toacher..tttenttion_scaaa
LogUnifora Uniform
10
10
fixed-teacher
n_inner
DiscreteUniform
16
32
meta
meta_learning_rate
LogUniform
10-6
meta
meta_moment
LogUnifort
10
10
meta
valication_split
Log Uniform
10
10
meta
Hyperparameter ranges the fixed-teacher method.
The parameters for the initial teacher weights are used exclusively by
Note that while the meta-lcarning approach has more hyperparameters than the baselines, the overall computational budget given each method is the same more hyperparameters are advantageous by default.
We run a hyperparameter search over the specified range for each combination of (method, n-training- examples). As hyperparameter optimization algorithm we use Tree-structured Parzen Estimator (Bergstra et al. 2011) for each method.
19


Reviews:
[{"review_id": "vI378HWJ0fQ", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nExperimental results on both synthetic and simulated datasets confirm the effectiveness of the proposed framework.  \n\nPros:\n1. The paper studies an important problem. Predicting the behavior of a dynamic system has many applications. \n\n2. The proposed model is interesting and may lead to a series of follow-up studies that leverage the strengths of both model-free and model-based methods using knowledge distillation techniques. \n\nCons:\n1. The baseline models are quite simple. There are stronger baselines as noted by the authors. While the proposal is a learning framework, it might still be worth customizing and comparing it with state-of-the-art models in specific problems. \n\n2. The presentation of the paper can be improved. It would be good to add a running example to explain the various concepts and definitions used in the paper. \n\nAdditional comments:\nTypo: \"a teacher networks\" => \"a teacher network\"; \"using only using\" => \"using only\"\n\n**Update after author response:** I appreciate the authors' efforts to address my comments. The new version reads better. However, I am still not entirely convinced by the choice of the simple baselines. Since a positive rating is already given, I would keep it unchanged. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "Gr9x8mxy_r", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a teacher-student training scheme to incorporate the useful information of trajectory to improve the predictive performance of model-free methods. The teacher network tries to \"guide\" the student network at the training stage by presenting an interpretation of the trajectory. The guidance is implemented by adding to the loss function a regularization term that penalizes the \"distance\" between the teacher's output and the hidden states of the student. The proposed method was tested and compared to other model-free methods.\n\nThe study in this work is interesting and important in RL. It tries to tackle the weakness of model-free methods by introducing the dynamics while avoiding to build a full model like the model-based methods. Searching for an optimal tradeoff between the two would benefit the practical uses.\n\nI'd vote for accepting the manuscript if the authors could address my concerns.\n\n- As mentioned in the text, the student internal h and the supervision signal h* are not required to have the same dimensionality. Are they required to have the same number of layers or certain correspondence between layers? Why or Why not? Do the authors have general principles on the design of the teacher and the teaching loss?\n\n- Is there any particular reason for using classification loss in all examples especially for a typically regression problem? Did regression loss lead to bad performance?\n\n- In the toy XOR example, the particular design of x* seems to play the key role. It is not the typical trajectory of x. This hand-crafted data does not help with demonstrating the teacher learning the powerful interpretation itself.\n\n- Though the authors have talked about why not to compare to a model-based method, I do not think it is convincing. As mentioned in the intro section, the model-based methods fail due to partial observability and etc, but I do not see the examples in this study have such issues. The computation of model-based methods depends on the complexity of internal model rather than the task. The argument should then be if the proposed method outperform model-based methods given the same complexity (e.g. the number of parameters) or same amount of data.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "8RThwa-6YXp", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a student-teacher framework, where the teacher network can be used to select and prioritize the relevant properties of the given dynamical system that should be learned by the student.\n\nPros:\n- (significance) I think the presented framework is a powerful one, and has a potential to be applied broadly to many real-world problems. \n- (quality) The development of the method is sound and well-motivated. The method was tested on a toy example and then applied to tasks with varying degrees of challenges.\n\nCons: (mostly on clarity)\n- What confuses me the most is the way the authors frames their work. From the current title \"Learning to interpret trajectories\", and the abstract, it was not really clear to me what the paper is about; I am not sure if the population of people who stop at the title would be the same as the population that finds the contents most interesting. Specifically:\n    - Why is \"trajectory\" a central keyword for this work? The word trajectory can be used in many different contexts, I don't think it was made clear anywhere in the paper what is the defining features of the \"trajectory-ness\" that the authors want to emphasize.\n    - What do you mean by \"interpret\"? This word is being used in a very loose fashion without a clear context; it is empty at best, and misleading at worst.\n- I get that the proposed framework avoids the *problems* of model-based and model-free methods, but I am having difficulties identifying what *advantages* of the two methods that the framework is incorporating.\n- One of the central questions that is raised in the introduction is this: \"What is the right way to incorporate information from different trajectories?\". But I am not sure how this work solves the problem of incorporating *different* trajectories specifically.\n\nAdditional comment:\n- The core concepts from the previous works that the work is based on, such as *learning using privileged information* or the meta-gradient approach, are not clearly introduced. Even a brief, ~1 sentence description would be helpful.\n- The ideas of model-based and model-free methods are reinforcement learning concepts, and may not be clear to people who are not in RL. Again some brief description would help.\n\nOverall, I have a mixed feeling about this paper and I currently stand between scores 5 and 6. Whereas I find the proposed method interesting, I feel that the lack of clarity, and the confounding of messages in the framing, make this paper rather short of the standard for the conference. \n\n**UPDATE:** \nMy major concerns were addressed in the revised version of the paper.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "-nyOOduxT-h", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis work tries to find a compromise of model-based and model-free methods, using a teacher and student network . The teacher network is trained with meta-gradients. It interprets the trajectories and provides activations for a student network that is supervised for a given task using the current state. \n\nStrengths:\n+ The paper is well motivated and solves interesting problem.\n+ The related work is thoroughly reviewed. \n\nWeaknesses:\n- Some claims made by authors are not validated. I suggest to add relevant citations in Sec.1. These claims support the motivation of this work but are not acceptable without proper references. For example, where is the evidence of deterioration on the tasks that are potentially relevant? Which task do model-free methods achieve substantially better performance?\n- The evaluation is conducted using the self-generated baselines. Why don't you use existing methods shown in Sec.2 and compare? I cannot find the result of the model-based baseline (as a counterpart of MF).\n- Typos: using only using model-free methods", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a new teacher-student framework where the teacher network guides the student network in learning useful information from trajectories of a dynamical system. The proposed framework is inspired by the Knowledge Distillation method. The teacher learns what information should be used from the trajectories and distills this information for the student in the form of target activations. In a nutshell, the framework allows the student to interpolate between model-based and model-free approaches in an automated fashion. Experimental evaluation on both the hand-crafted and simulated tasks demonstrate the effectiveness of the proposed framework. The reviewers had borderline scores in their initial reviews and raised several questions for the authors. The reviewers appreciated the rebuttal, which helped in answering their key questions -- I want to thank the authors for engaging with the reviewers during the discussion phase. The reviewers have an overall positive assessment of the paper, and believe that the proposed teacher-student framework is novel and potentially useful for many real-world problems. The reviewers have provided detailed feedback in their reviews, and I would like to strongly encourage the authors to incorporate this feedback when preparing the final version of the paper.
Decision: Accept (Poster)