On Lottery Tickets and Minimal Task Representations in Deep Reinforcement Learning

The lottery ticket hypothesis questions the role of overparameterization in supervised deep learning . But how is the performance of winning lottery tickets affected by the distributional shift inherent to reinforcement learning problems ? In this work , we address this question by comparing sparse agents who have to address the non-stationarity of the exploration-exploitation problem with supervised agents trained to imitate an expert . We show that feed-forward networks trained with behavioural cloning compared to reinforcement learning can be pruned to higher levels of sparsity without performance degradation . This suggests that in order to solve the RL problem agents require more degrees of freedom . Using a set of carefully designed baseline conditions , we find that the majority of the lottery ticket effect in both learning paradigms can be attributed to the identified mask rather than the weight initialization . The input layer mask selectively prunes entire input dimensions that turn out to be irrelevant for the task at hand . At a moderate level of sparsity the mask identified by iterative magnitude pruning yields minimal task-relevant representations , i.e. , an interpretable inductive bias . Finally , we propose a simple initialization rescaling which promotes the robust identification of sparse task representations in low-dimensional control tasks . 1 INTRODUCTION . Recent research on the lottery ticket hypothesis ( LTH , Frankle & Carbin , 2019 ; Frankle et al. , 2019 ) in deep learning has demonstrated the existence of very sparse neural networks that train to performance levels comparable to those of their dense counterparts . These results challenge the role of overparameterization in supervised learning and provide a new perspective on the emergence of stable learning dynamics ( Frankle et al. , 2020a ; b ) . Recently these results have been extended to various domains beyond supervised image classification . These include self-supervised learning ( Chen et al. , 2020a ) , natural language processing ( Yu et al. , 2019 ; Chen et al. , 2020b ) and semantic segmentation ( Girish et al. , 2020 ) . But how does the lottery ticket ticket phenomenon transfer to reinforcement learning agents ? One key challenge may be the inherent non-stationarity of the optimization problem in deep reinforcement learning ( DRL ) : The data-generation process is not static , but depends on the changing state of the neural network . Furthermore , a weight may serve different roles at different stages of learning ( e.g . during exploration and exploitation ) . It is not obvious how a simple weight magnitude-based pruning heuristic acting on a well-performing policy shapes the learning process of the agent . In this work , we therefore investigate winning tickets in reinforcement learning and their underlying contributing factors . We compare supervised behavioral cloning with DRL , putting a special emphasis on the resulting input representations used for prediction and control . Thereby , we connect the statistical perspective of sparse structure discovery ( e.g . Hastie et al. , 2019 ) with the iterative magnitude pruning ( IMP , Han et al. , 2015 ) procedure in the context of Markov decision processes ( MDPs ) . The contributions of this work are summarized as follows : 1 . We show that winning tickets exist in both high-dimensional visual and control tasks ( continuous/discrete ) . A positive lottery ticket effect is robustly observed for both off-policy DRL algorithms , including Deep-Q-Networks ( DQN , Mnih et al. , 2015 ) and on-policy policy-gradient methods ( PPO , Schulman et al. , 2015 ; 2017 ) , providing evidence that the lottery ticket effect is a universal phenomenon across optimization formulations in DRL . 2 . By comparing RL to supervised behavioral cloning ( BC ) , we show that networks trained with explicit supervision can be pruned to higher sparsity levels before performance starts to degrade , indicating that the RL problem requires a larger amount of parameters to address exploration , distribution shift & quality of the credit assignment signal ( section 3.1 ) . 3 . By introducing a set of lottery ticket baselines ( section 3 ; figure 1 , middle column ) , we disentangle the contributions of the mask , weight initialization and layer-wise pruning ratio . We demonstrate that the mask explains most of the ticket effect in behavioral cloning and reinforcement learning for MLP-based agents , whereas the associated weight initialization is less important ( section 3.2 and 3.3 ) . For CNN-based agents the weight initialization contributes more . 4 . By visualizing the sparsified weights for each layer , we find that early network layers are pruned more . Entire input dimensions can be rendered invisible to MLP-based agents by the pruning procedure . By this mechanism , IMP compresses the input representation of the MDP ( e.g . figure 1 , left column , bottom row ) and reveals a minimal task representation for the underlying control problems ( section 4 ) . 5 . The IMP input layer mask not only eliminates obviously redundant dimensions , but also identifies complex relationships between input features and the task of the agent ( section 4.1 ) , e.g . the proximity of an enemy or the speed of an approaching object . We show that the input masking can be transferred to train dense agents with sparse inputs at lower costs . 6 . We show that the weight initialization scheme is important for discovering minimal representations . Depending on the input size of different layers of the network , global magnitude-based pruning can introduce a strong layer-specific pruning bias . We compare initializations and show that a suitable initialization scheme enables the removal of task-irrelevant dimensions ( section 4.2 ) . 2 BACKGROUND AND RELATED WORK . Iterative Magnitude Pruning . We use the iterative pruning procedure outlined in Frankle & Carbin ( 2019 ) to identify winning tickets . We train DRL agents for a previously calibrated number of transitions and track the best performing network checkpoint throughout . Performance is measured by the average return on a set of evaluation episodes . Afterwards , we prune 20 % of the weights with smallest magnitude globally ( across all layers ) . The remaining weights are reset to their initial values and we iterate this procedure ( train→ prune→ reset ) .1 The lottery ticket effect refers to the performance gap between the sparse network obtained via IMP and a randomly initialized network with sparsity-matched random pruning mask . Lottery Tickets in Deep Reinforcement Learning . Yu et al . ( 2019 ) previously demonstrated the existence of tickets in DRL that outperform parameter-matched random initializations . They obtained tickets for a distributed on-policy actor-critic agent on a subset of environments in the ALE benchmark ( Bellemare et al. , 2013 ) as well as a set of discrete control tasks . While they provide empirical evidence for the existence of lottery tickets in DRL , they did not investigate the underlying mechanisms . Here , we aim to unravel these mechanisms . To this end , we focus on a diverse set of environments and provide a detailed comparison between supervised behavioral cloning and on-/offpolicy Deep RL with a set of carefully designed ticket baselines . We analyze the resulting masked representations that the agent learns to act upon and the impact of specific weight initializations on the resulting sparse networks . Lottery Tickets with Non-Stationary Data Distributions . Desai et al . ( 2019 ) investigated whether trained lottery tickets overfit the training data distribution under which they were obtained . Using transfer learning tasks on natural language data , they showed that lottery tickets provide general inductive biases . Similar ticket transfer results were reported by Morcos et al . ( 2019 ) and Mehta ( 2019 ) in the context of optimizers and vision datasets . Unlike our work , these studies do not investigate within-training covariate shift , but instead focus on transferring ticket initializations after a full IMP run . Chen et al . ( 2021 ) , on the other hand , investigate the ticket phenomenon in the context of lifelong learning and class-incremental image classification . They propose new pruning strategies to overcome the sequential nature of tasks and need for increased model capacity . Compared to the DRL setting , the covariate shift is here determined by the curriculum schedule of tasks and not the exploration behaviour of the network-parameterized agent . Deep Reinforcement Learning Background . In our off-policy DRL experiments , we train DeepQ-Networks ( DQN , Mnih et al. , 2015 ) with double Q-learning loss ( Van Hasselt et al. , 2016 ) and prioritized experience replay ( Schaul et al. , 2015 ) . As a representative on-policy algorithm , we chose Proximal Policy Optimization ( PPO , Schulman et al. , 2015 ; 2017 ) . PPO is a baseline-corrected policy gradient algorithm which uses a clipping strategy to approximate a computationally expensive trust-region optimization method . For illustrative purposes , we train DQN agents on a visual navigation task , in which an agent has to collect coins in a grid while avoiding poison and two patrollers that are moving in restricted parts of the grid ( figure 1 , left column , bottom row ; SI B ) . We scale our results to four PyBullet ( Ellenberger , 2018 ) continuous control and a subset of ALE benchmark ( Bellemare et al. , 2013 ) environments . Due to computational considerations we limit each individual IMP iteration for the ATARI environments to 2.5 million frames . All other tasks were trained for a pre-calibrated generous amount of transitions . We focus on feedforward value estimators and policies ( MLP & CNN ) and used default hyperparameters with little tuning ( SI C ) . Supervised Behavioral Cloning . While most supervised learning relies on a stationary data distribution provided by a static dataset , reinforcement learning agents have to acquire their training data in an action-perception loop . Since the agent ’ s behavioural policy is learned over time , the data distribution used in optimization undergoes covariate shift . To study how the covariate shift , additional exploration problem and different credit assignment signal influence winning tickets , we mimic the supervised learning case by training agents via supervised policy distillation ( Rusu et al. , 2015 ; Schmitt et al. , 2018 ) . We roll out a pre-trained expert policy and train the student agent by minimizing the KL divergence between the student ’ s and teacher ’ s policies . 1In supervised learning , the pruning mask is often constructed based on an early stopping criterion and the final network . We instead track the best performing agent . Thereby , we reduce noise introduced by unstable learning dynamics and exploit that the agent is trained and evaluated on the same environment . We found that late rewinding to a later checkpoint ( Frankle et al. , 2019 ) is not necessary for obtaining tickets ( SI figure 14 ) . 3 DISENTANGLING TICKET CONTRIBUTIONS IN BC AND DEEP RL . There are two contributing factors to the lottery ticket effect : The IMP-identified binary mask and the preserved initialized weights that remain after pruning ( mask/weights ) . We aim to disentangle the contributions by introducing a set of counterfactual baselines , which modify the original IMP procedure ( figure 1 , middle column ; table 1 ) . A first baseline estimates how much of the performance of the ticket can be attributed to the initial weights , by means of a layer-specific permutation of the weights that remain after masking ( mask/permuted ) . A second , weaker baseline estimates the contribution of the mask , by also permuting the layer-specific masks ( permuted/permuted ) . Finally , we consider the standard random/re-init baseline , which samples random binary masks – discarding layer-specific pruning ratios – and re-initializes all weights at each IMP iteration . Throughout the next sections we use these baselines to analyze and compare the factors that give rise to the lottery ticket effect in different control settings .

Summary: This paper investigates the Lottery Ticket hypothesis in the context of deep RL for identification of sparse task representation in low-dimensional control tasks. This is primarily an empirical investigation where several experiments and consequent analysis reveal what a "winning ticket" means in case of policy models or Q function approximators given that Deep RL is prone to gradual distribution shift. The experiments also help to identify how standard magnitude pruning will behave under different environment updates, feedback schedules or initialization dynamics.