The Role of Pretrained Representations for the OOD Generalization of Reinforcement Learning Agents
This work addresses the problem of sample-efficient OOD generalization for reinforcement learning agents, particularly in robotics, but it is incremental as it builds on existing representation learning approaches.
The study investigated how pretrained VAE-based representations affect out-of-distribution generalization in reinforcement learning agents, finding that many agents are robust to realistic distribution shifts and that a simple proxy task can reliably predict generalization performance.
Building sample-efficient agents that generalize out-of-distribution (OOD) in real-world settings remains a fundamental unsolved problem on the path towards achieving higher-level cognition. One particularly promising approach is to begin with low-dimensional, pretrained representations of our world, which should facilitate efficient downstream learning and generalization. By training 240 representations and over 10,000 reinforcement learning (RL) policies on a simulated robotic setup, we evaluate to what extent different properties of pretrained VAE-based representations affect the OOD generalization of downstream agents. We observe that many agents are surprisingly robust to realistic distribution shifts, including the challenging sim-to-real case. In addition, we find that the generalization performance of a simple downstream proxy task reliably predicts the generalization performance of our RL agents under a wide range of OOD settings. Such proxy tasks can thus be used to select pretrained representations that will lead to agents that generalize.