Transient Non-stationarity and Generalisation in Deep Reinforcement Learning

1 INTRODUCTION . In RL , as an agent explores more of its environment and updates its policy and value function , the data distribution it uses for training changes . In deep RL , this non-stationarity is often not addressed explicitly . Typically , a single neural network model is initialised and continually updated during training . Conventional wisdom about catastrophic forgetting ( Kemker et al. , 2018 ) implies that old updates from a different data-distribution will simply be forgotten . However , we provide evidence for an alternative hypothesis : networks exhibit a memory effect in their learned representations which can harm generalisation permanently if the data-distribution changed over the course of training . To build intuition , we first study this phenomenon in a supervised setting on the CIFAR-10 dataset . We artificially introduce transient non-stationarity into the training data and investigate how this affects the asymptotic performance under the final , stationary data in the later epochs of training . Interestingly , we find that while asymptotic training performance is nearly unaffected , test performance degrades considerably , even after the data-distribution has converged . In other words , we find that latent representations in deep networks learned under certain types of non-stationary data can be inadequate for good generalisation and might not be improved by later training on stationary data . Such transient non-stationarity is typical in RL . Consequently , we argue that this observed degradation of generalisation might contribute to the inferior generalisation properties recently attributed to many RL agents evaluated on held out test environments ( Zhang et al. , 2018a ; b ; Zhao et al. , 2019 ) . Furthermore , in contrast to Supervised Learning ( SL ) , simply re-training the agent from scratch once the data-distribution has changed is infeasible in RL as current state of the art algorithms require data close to the on-policy distribution , even for off-policy algorithms like Q-learning ( Fedus et al. , 2020 ) . To improve generalisation of RL agents despite this restriction , we propose Iterated Relearning ( ITER ) . In this paradigm for deep RL training , the agent ’ s policy and value are periodically distilled into a freshly initialised student , which subsequently replaces the teacher for further optimisation . While this occasional distillation step simply aims to re-learn and replace the current policy and ∗University of Oxford . Corresponding author : Maximilian Igl ( maximilian.igl @ gmail.com ) †Now at DeepMind , London ‡Delft University of Technology value outputs for the training data , it allows the student to learn a better latent representation with improved performance for unseen inputs because it eliminates non-stationarity during distillation . We propose a practical implementation of ITER which performs the distillation in parallel to the training process without requiring additional training data . While this introduces a small amount of non-stationarity into the distillation step , it greatly improves sample efficiency without noticeably impacting performance . Experimentally , we evaluate ITER on the Multiroom environment , as well as several environments from the recently proposed ProcGen benchmark and find that it improves generalisation . This provides further support to our hypothesis and indicates that the non-stationarity inherent to many RL algorithms , even when training on stationary environments , should not be ignored when aiming to learn robust agents . Lastly , to further support this claim and provide more insight into possible causes of the discovered effect , we perform additional ablation studies on the CIFAR-10 dataset . 2 BACKGROUND . We describe an RL problem as a Markov decision process ( MDP ) ( S , A , T , r , p0 , γ ) ( Puterman , 2014 ) with actions a∈A , states s∈S , initial state s0∼p0 , transition dynamics s′∼T ( s , a ) , reward function r ( s , a ) ∈R and discount factor γ . The unnormalised discounted state distribution induced by a policy π is defined as dπ ( s ) = ∑∞ t=0 γ tPr ( St=s|S0∼p0 , At∼π ( ·|St ) , St+1∼T ( St , At ) ) . In ITER , we learn a sequence of policies and value functions , which we denote with π ( k ) ( a|s ) and V ( k ) ( s ) at the kth iteration ( k ∈ { 0 , 1 , 2 , . . . } ) , parameterized by θk . We briefly discuss some forms of non-stationarity which can arise in RL , even when the environment is stationary . For simplicity , we focus the exposition on actor-critic methods which use samples from interaction with the environment to estimate the policy gradient given by g = E [ ∇θ log πθ ( a|s ) Aπ ( s , a , s′ ) |s , a , s′ ∼ dπ ( s ) π ( a|s ) T ( s′|s , a ) ] . The advantage is often estimated as Aπ ( s , a , s′ ) = r ( s , a ) + γV π ( s′ ) − V π ( s ) . Typically , we also use neural networks to approximate the baseline V πφ ( s ) and for bootstrapping from the future value V π φ ( s ′ ) . φ can be learned by minimising E [ Aπ ( s , a , s′ ) 2 ] by stochastic semi-gradient descent , treating V πφ ( s′ ) as a constant . There are at least three main types of non-stationarity in deep RL . First , we update the policy πθ , which leads to changes in the state distribution dπθ ( s ) . Early on in training , a random policy πθ only explores states close to initial states s0 . As πθ improves , new states further from s0 are encountered . Second , changes to the policy also change the true value function V π ( s ) which V πφ ( s ) is approximating . Lastly , due to the use of bootstrap targets in temporal difference learning , the learned value V πφ ( s ) is not regressed directly towards V π ( s ) . Instead V πφ fits a gradually evolving target sequence even under a fixed policy π , thereby also changing the policy gradient estimator g . 3 THE IMPACT OF NON-STATIONARITY ON GENERALISATION . In this section we investigate how asymptotic performance is affected by changes to the datadistribution during training . In particular , we assume an initial , transient phase of non-stationarity , followed by an extended phase of training on a stationary data-distribution . This is similar to the situation in RL where the data-distribution is affected by a policy which converges over time . We show that this transient non-stationarity has a permanent effect on the learned representation and negatively impacts generalisation . As interventions in RL training can lead to confounding factors due to off-policy data or changed exploration behaviour , we utilise Supervised Learning ( SL ) here to provide initial evidence in a more controlled setup . We use the CIFAR-10 dataset for image classification ( Krizhevsky et al. , 2009 ) and artificially inject non-stationarity . Our goal is to provide qualitative results on the impact of non-stationarity , not to obtain optimal performance . We use a ResNet18 ( He et al. , 2016 ) architecture , similar to those used by Espeholt et al . ( 2018 ) and Cobbe et al . ( 2019a ) . Parameters are updated using Stochastic Gradient Descent ( SGD ) with momentum and , following standard practice in RL , we use a constant learning rate and do not use batch normalisation . Weight decay is used for regularisation . Hyper-parameters and more details can be found in appendix B . We train for a total of 2500 epochs . While the last 1500 epochs are trained on the full , unaltered dataset , we modify the training data in three different ways during the first 1000 epochs . Test data is left unmodified throughout training . While each modification is loosely motivated by the RL setting , our goal is not to mimic it exactly ( which would be infeasible ) , nor to disentangle the contributions of different types of non-stationarity . Instead , we aim to show that these effects reliably occur in the presence to various types of non-stationarity , and provide intuitions that can be brought into the RL setting in Section 4 . For the first modification , called Dataset Size , we initially train only on a small fraction of the full dataset and gradually add more data points after each epoch , at a constant rate , until the full dataset is available after epoch 1000 . During the non-stationary phase , data points are reused multiple times to ensure the same number of network updates are made in each epoch . For the modification Wrong Labels we replace all training labels by randomly drawn incorrect ones . After each epoch , a constant number of these labels are replaced by their correct values . Lastly , Noisy Labels is similar to Wrong Labels , but the incorrect labels are sampled uniformly at the start of each epoch . For both , all training labels are correct after epoch 1000 . While Dataset Size is inspired by the changing state distribution seen by an evolving policy , Wrong Labels and Noisy Labels are motivated by the consistent bias or fluctuating errors a learned critic can introduce in the policy gradient estimate . The results are shown in fig . 1 . While the final training accuracy ( left ) is almost unaffected ( see table 1 in the appendix for exact results ) , all three nonstationary modifications significantly reduce the test accuracy ( right ) . The plateau in accuracy shows that this effect persists even after the models are further trained using the full dataset with correct labels : non-stationarity early in training has a permanent effect on the learned representations and quality of generalisation . These results indicate that the non-stationarity introduced by the gradual convergence of the policy in RL might similarly de- teriorate the generalisation of the agent . To overcome this , we propose ITER in the next section . The key insight enabling ITER is that the observed negative effect is restricted to the test data , whereas the predictions on the training data are unaffected and of high quality . 4 ITER . In section 3 , we have seen evidence that the non-stationarity which is present in many deep RL algorithms might lead to impaired generalisation on held-out test environments . To mitigate this and improve generalisation to previously unseen states , we propose Iterated Relearning ( ITER ) : instead of updating a single agent model throughout the entire training process , ITER learns a sequence of models , each of which is exposed to less non-stationarity during its training . As we will show in section 5 , this improves generalisation . ITER can be applied on top of a wide range of base RL training methods . For simplicity , we focus in the following exposition on actor-critic methods and use Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) for the experimental evaluation . The underlying insight behind ITER is that at any point during RL training the latent representation of our current agent network might be significantly damaged by non-stationarity , but its outputs on the training data are comparatively unaffected ( see fig . 1 ) . Consequently , ITER aims to periodically replace the current agent network , the ‘ teacher ’ , by a ‘ student ’ network which was freshly initialised and trained to mimic the teacher on the current training data . Because this re-learning and replacement step can be performed on stationary data , it allows us to re-learn a policy that matches performance on the training data but generalises better to novel test environments . ITER begins with an initial policy π ( 0 ) and value function V ( 0 ) and then repeats the following steps , starting with iteration k = 0 . 1 . Use the base RL algorithm to train π ( k ) and V ( k ) . 2 . Initialise new student networks for π ( k+1 ) and V ( k+1 ) . We refer to the current policy π ( k ) and value function V ( k ) as the teacher . 3 . Distill the teacher into the student . This phase is discussed in more detail in section 4.1 . 4 . The student replaces the teacher : π ( k ) and V ( k ) can be discarded . 5 . Increment k and return to step 1 . Repeat as many times as needed . This results in alternating RL training with distillation into a freshly initialised student . The RL training phases continue to introduce non-stationarity until the models converge , so we want to iterate the process , repeating steps 1-4 . How often we do so depends on the environment and can be chosen as a hyper-parameter . In practise we found the results to be quite robust to this choice and recommend , as general rule , to iterate as often as possible within the limits outlined in section 4.2 . There , we introduce a practical implementation of ITER which re-uses data between steps 1 and 3 in order to not require additional samples from the environment .

Summary: The paper deals with a relevant issue. The simplified supervised learning setting is a good way of looking at the issue of non-stationarity in isolation and it makes a compelling case that neural networks optimized by SGD can have generalization issues in settings where the data distribution changes over time, even after the data distribution converges. The solution proposed by the paper is simple and can be applied to most off-the-shelf RL algorithms. My main criticism would be that the hybrid objective feels rather ad-hoc and that the proposed method could use a bit more theoretical justification. Also, since the issue being tackled here is quite general to RL, a wider set of benchmarks and base-algorithms (in addition to PPO) would be necessary to get a better picture.