SiT: Symmetry-Invariant Transformers for Generalisation in Reinforcement Learning

TL;DR

SiT introduces a symmetry-aware vision transformer that enhances generalization and sample efficiency in reinforcement learning by preserving graph symmetries through a novel attention mechanism.

Contribution

The paper proposes Graph Symmetric Attention within a scalable ViT, enabling invariance to symmetries and improving RL generalization and sample efficiency.

Findings

Outperforms ViTs on MiniGrid and Procgen benchmarks

Achieves better sample efficiency on Atari 100k and CIFAR10

Demonstrates improved generalization in RL environments

Abstract

An open challenge in reinforcement learning (RL) is the effective deployment of a trained policy to new or slightly different situations as well as semantically-similar environments. We introduce Symmetry-Invariant Transformer (SiT), a scalable vision transformer (ViT) that leverages both local and global data patterns in a self-supervised manner to improve generalisation. Central to our approach is Graph Symmetric Attention, which refines the traditional self-attention mechanism to preserve graph symmetries, resulting in invariant and equivariant latent representations. We showcase SiT's superior generalization over ViTs on MiniGrid and Procgen RL benchmarks, and its sample efficiency on Atari 100k and CIFAR10.