Transient Non-Stationarity and Generalisation in Deep Reinforcement Learning

TL;DR

This paper identifies transient non-stationarity as a key factor harming generalisation in deep RL and demonstrates that neural networks retain legacy features that can impair performance on unseen states. It introduces Iterated Relearning (ITER), a distillation-based framework that periodically retrains a freshly initialised student to reduce non-stationarity, combining standard RL training with distillation losses. Empirical results across Multiroom, Boxoban, ProcGen, and CIFAR-10 illustrate improved generalisation and sample efficiency, supporting the idea that mitigating non-stationarity yields more robust representations. The work offers a practical, parallelisable approach and a plausible mechanism—the legacy feature hypothesis—for why transient non-stationarity degrades generalisation in deep RL.

Abstract

Non-stationarity can arise in Reinforcement Learning (RL) even in stationary environments. For example, most RL algorithms collect new data throughout training, using a non-stationary behaviour policy. Due to the transience of this non-stationarity, it is often not explicitly addressed in deep RL and a single neural network is continually updated. However, we find evidence that neural networks exhibit a memory effect where these transient non-stationarities can permanently impact the latent representation and adversely affect generalisation performance. Consequently, to improve generalisation of deep RL agents, we propose Iterated Relearning (ITER). ITER augments standard RL training by repeated knowledge transfer of the current policy into a freshly initialised network, which thereby experiences less non-stationarity during training. Experimentally, we show that ITER improves performance on the challenging generalisation benchmarks ProcGen and Multiroom.

Figures (9)

• Figure 1: Accuracy on CIFAR-10 when the training data is non-stationary over the first 1000 epochs (dashed line). The remaining epochs are trained on the full, unaltered training data. Testing is performed on unaltered data throughout. While final training performance (left) is almost unaffected, test accuracy (right) is significantly reduced by initial, transient non-stationarity.
• Figure 2: Evaluation on Multiroom and Boxoban. Shown are mean and standard error over twelve seeds. Left: Return for with and without on Multiroom. Dotted lines indicate when the network was replaced by a new student. Middle: Evaluation on layouts with a fixed number of rooms; training is still with a random number of rooms. 's advantage is more pronounced for harder levels. Right: Return on Boxoban.
• Figure 3: Evaluation on ProcGen. Dashed lines indicate replacing the teacher. Left: Test performance averaged over six environments (StarPilot, Dodgeball, Climber, Ninja, Fruitbot and BigFish). Shown are mean and standard error over all 30 runs (five per environment). Results are normalised by the final test-performance of the baseline on each respective environment to make them comparable. We also compare against the previous state of the art method IBAC-SNIigl2019generalization. Middle: Evaluation on Climber. improves test performance without improving training, supporting our claim that improves the latent representation of the agent. Right: Evaluation on BigFish. On some environments, improves both train- and test- performance.
• Figure 4: Left: Ablation studies with sequential and without terms $\mathcal{L}_{\text{PG}}$ and $\mathcal{L}_{\text{TD}}$ (\ref{['eq:pg']}). Right: Schematic depiction of training setup for \ref{['fig:sl']} (middle and right). More details are given in \ref{['sec:exp:sl']}. $\mathcal{D}$ is the unmodified CIFAR-10 training data while for $\mathcal{D}^{f,m}$ modification $m\in\{\texttt{Noisy Labels},\texttt{Wrong Labels}, \texttt{Dataset Size}\}$ is applied to the fraction $(1-f)$ of all data-points. In this two phase training setup, we first train on $\mathcal{D}^{f,m}$ during phase 1 and continue on $\mathcal{D}$ during phase 2. A linear predictor parameterised by $(W^{f,m}_i,b^{f,m}_i)$ is trained on $\mathcal{D}$ after each phase $i$, while holding the encoder $\phi_i^{f,m}(x)$ fixed. Evaluation of the resulting classifiers is performed on the original test data. Classifier $i=1$ measures the relevance of the legacy features while classifier $i=2$ measure the final generalisation performance.
• Figure 5: Left: Test accuracy of students (solid lines) that only learn to mimic the behaviour of poorly generalising teachers in \ref{['fig:sl-nonstat']} (dashed lines). Middle: Final test accuracy of networks trained consecutively on two different datasets. The $x$-axis shows the accuracy of using encoders trained on the first dataset, retraining only the last layer on the second: nearly useless earlier representations impact future learning much less than slightly sub-optimal ones. Markers indicate modifications to first dataset; colours indicate the fraction of unmodified data points $f$. Dashed line shows accuracy for $f=1$. Right: Singular values of feature matrix $\Phi$, normalised by the largest SV. Solid lines show intermediate values of $f$ with low test accuracy, dashed lines small values of $f$ with higher accuracy. More plots can be found in the appendix.