Improving Zero-Shot Generalization in Reinforcement Learning Through Abstract Representations

Reinforcement learning algorithms suffer from zero-shot generalization, where a learned policy is tested in new, out-of-distribution environments. One of the key reasons is the overfitting of the representation learning model. Previous studies have used data augmentation or contrastive learning to address this overfitting. Yet these methods often rely on task-specific augmentations or hand-picked negative samples, which are not robust to various tasks. We hypothesize that learning more abstract representations by ignoring pixel-level details allows the representation learning model to generalize better. Moreover, when an agent encounters a scene similar to one seen in the training environment but with slight variations in the background, it will ignore the differences and treat the scene as invariant, making the correct decision. Based on this idea, we introduce Dream with Abstractions (Dr. Abs), an algorithm that requires no custom augmentations or negative samples. First, we train a representation network with a joint embedding predictive architecture, where it learns an abstract representation. Next, we use the learned representation to regularize the latent states inside the reinforcement-learning agent. By ignoring pixel-level details, our method outperforms the baseline algorithm by approximately 22% on the widely used zero-shot generalization setup, thereby enhancing its generalization ability. Furthermore, Dr. Abs can be used in conjunction with other RL methods and can work in tandem with techniques such as data augmentation.

Authors

• Tsai Shuo Kuo
• Ryutaro Ichise

Keywords

No keywords are indexed for this paper.