Sparse Masked Attention Policies for Reliable Generalization

TL;DR

This paper introduces a learned masking approach within attention-based policy networks in reinforcement learning, significantly enhancing the generalization of policies to unseen tasks by reliably removing unnecessary information.

Contribution

It proposes a novel information removal method using a learned masking function integrated with attention weights, improving policy generalization in unseen environments.

Findings

Significant improvement in generalization to unseen tasks in Procgen benchmark

Outperforms standard PPO and masking approaches

Enhances reliability of information removal for better policy transfer

Abstract

In reinforcement learning, abstraction methods that remove unnecessary information from the observation are commonly used to learn policies which generalize better to unseen tasks. However, these methods often overlook a crucial weakness: the function which extracts the reduced-information representation has unknown generalization ability in unseen observations. In this paper, we address this problem by presenting an information removal method which more reliably generalizes to new states. We accomplish this by using a learned masking function which operates on, and is integrated with, the attention weights within an attention-based policy network. We demonstrate that our method significantly improves policy generalization to unseen tasks in the Procgen benchmark compared to standard PPO and masking approaches.