Neural Distillation as a State Representation Bottleneck in Reinforcement Learning

Despite the impressive successes of modern reinforcement learning (RL) (Sutton & Barto, 2018) methods in designing efficient specialized control policies for a wide variety of difficult tasks, many studies have highlighted the limited ability of RL agents to generalize to variations of such tasks that would appear easy to a human being (Farebrother et al., 2018; Packer et al., 2018; Zhang et al., 2018; Song et al., 2020; Cobbe et al., 2019). This work is motivated by the idea that networks trained for specific tasks build state representations that can easily be fooled by the ambiguity between observation variables. For instance, in some platform video games, it is possible to design an optimal policy for a given level based solely on background features and progression indicators, rather than on the position of platforms and enemies (Song et al., 2020). While very efficient on this specific level, such a policy might not perform well on another. Conversely, we formulate and evaluate the hypothesis that a network trained to imitate several such specialized policies on a limited set of task variations induces a state representation that lifts the ambiguity and filters out confounding observation variables. Specifically, we investigate whether the process of network distillation (Hinton et al., 2015; Rusu et al., 2016a), inspired by knowledge consolidation in cognitive systems (Wilson & McNaughton, 1994; Ashworth et al., 2014; McClelland et al., 1995), induces valuable state representations. The interplay between distillation and state representation appears to have received little attention so far. We endeavor to fill this gap and investigate how neural distillation can act as a state representation bottleneck in RL. Our contributions are as follows: We propose a generic experimental protocol to evaluate the effects of imitation (via distillation) on state representation.