SiT: Symmetry-invariant Transformers for Generalisation in Reinforcement Learning

Matthias Weissenbacher, Rishabh Agarwal, Yoshinobu Kawahara

Primary Area: reinforcement learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: Reinforcement Learning, Generalization, Self-Attention, Equivariant Transformer, Invariant Transformer

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Abstract:

A major open challenge in reinforcement learning (RL) is the effective deployment of a trained policy to new and slightly different situations i.e. out-of-distribution data, as well as semantically-similar environments. To overcome these limitations, we introduce the **S**ymmetry-**I**nvariant **T**ransformer (**SiT**), a scalable invariant and equivariant Transformer model that identifies and utilizes local
and global patterns in the data. Central to our approach is the Graph Symmetric Attention (GSA) mechanism, adapting the self-attention mechanism to maintain graph symmetries. The invariant and equivariant latent representations are then used as a starting point for a subsequent policy and value networks. By striking a balance between local symmetries and overarching data trends, our model achieves innate capabilities to handle unfamiliar data distributions. We demonstrate improved generalization using our approach on MiniGrid and Procgen environments.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 4598