# SiT: Symmetry-Invariant Transformers for Generalisation in Reinforcement Learning

arXiv: https://arxiv.org/abs/2406.15025 · HTML: https://mlsanity.com/paper/2406.15025/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in RL for robot locomotion and manipulation, one of 15 directions in Deep reinforcement learning methods.
- Builds on 43 of its 70 references that are in the map.
- First posted on arXiv 21 Jun 2024.

## Placed in

- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

## Builds on (most cited first, 30 of 43)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Dynamic Routing Between Capsules](https://mlsanity.com/paper/1710.09829/) (2017)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [General E(2)-Equivariant Steerable CNNs](https://mlsanity.com/paper/1911.08251/) (2019)
- [Invariant and Equivariant Graph Networks](https://mlsanity.com/paper/1812.09902/) (2018)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Towards a Definition of Disentangled Representations](https://mlsanity.com/paper/1812.02230/) (2018)
- [Escaping the Big Data Paradigm with Compact Transformers](https://mlsanity.com/paper/2104.05704/) (2021)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning](https://mlsanity.com/paper/2006.16908/) (2020)
- [CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention](https://mlsanity.com/paper/2108.00154/) (2021)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)

## Built on by (most cited first, 3 of 3)

- [Approximate Equivariance in Reinforcement Learning](https://mlsanity.com/paper/2411.04225/) (2024)
- [DreamSAC: Learning Hamiltonian World Models via Symmetry Exploration](https://mlsanity.com/paper/2603.07545/) (2026)
- [Sparse Masked Attention Policies for Reliable Generalization](https://mlsanity.com/paper/2602.19956/) (2026)

## Abstract (arXiv)

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.