On Lottery Tickets and Minimal Task Representations in Deep Reinforcement Learning
This work provides insights into overparameterization and sparsity in RL, which could help improve efficiency and interpretability for AI researchers, though it is incremental in extending the lottery ticket hypothesis to RL.
The study investigated how the lottery ticket hypothesis applies to reinforcement learning (RL) by comparing sparse agents in RL and supervised imitation learning, finding that RL requires more parameters to handle distributional shift and that input layer pruning yields interpretable minimal task representations.
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-specific distributional shift 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.