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: This paper investigates an interesting problem that transient non-stationarity can affect the generalization of the neural network. This paper first conducts experiments on a supervised learning task to illustrate that transient non-stationarity can lead to degenerated performance on testing set. Then, the paper proposes an RL algorithm called ITER to avoid the negative impact of such non-stationarity.