SiT: Symmetry-Invariant Transformers for Generalisation in Reinforcement Learning

Matthias Weissenbacher, Rishabh Agarwal, Yoshinobu Kawahara
6/21/2024
cs.LG

No official code found. Best implementation: matthias-weissenbacher/SiT (0 stars, last commit 2024-06-21).

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.

View on arXivView PDF

Code Implementations(1)

matthias-weissenbacher/SiTCommunity100%

Symmetry-invariant Transformers for Generalisation in Reinforcement Learning

01PythonMay 28, 20242 years ago

Similar Papers

Multi-Group Equivariant Augmentation for Reinforcement Learning in Robot Manipulation

85% match2025

Partial Symmetry Enforced Attention Decomposition (PSEAD): A Group-Theoretic Framework for Equivariant Transformers in Biological Systems

85% match2025

Partially Equivariant Reinforcement Learning in Symmetry-Breaking Environments

83% match2025

Reinforcement Learning Using known Invariances

83% match2025

Vanilla Group Equivariant Vision Transformer: Simple and Effective

82% match2026

Symmetry Breaking in Transformers for Efficient and Interpretable Training

82% match1 repo12026

Cite this paper

@article{weissenbacher2024sit,
  title  = {SiT: Symmetry-Invariant Transformers for Generalisation in Reinforcement Learning},
  author = {Matthias Weissenbacher and Rishabh Agarwal and Yoshinobu Kawahara},
  year   = {2024},
  eprint = {2406.15025},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2406.15025v1},
  journal = {ICML 2024 2024}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2406.15025.svg)](https://codeofpaper.com/paper/2406.15025)

Discussion