Research map / Reinforcement learning / Visual and representation RL
SiT: Symmetry-invariant Transformers for Generalisation in Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SiT, a ViT-based policy architecture built on Graph Symmetric Attention (GSA) that embeds grid symmetries (rotation, flip, shift) into self-attention, and evaluates it for generalisation in RL (MiniGrid, Procgen) plus sample efficiency on Atari 100k, DM-control and CIFAR-10.
Problem
Policies trained in image-based RL often fail to generalise to new or slightly different situations, and standard ViTs are permutation invariant and perform poorly in RL with limited sample efficiency and hyperparameter sensitivity.
Approach
Replace standard self-attention with Graph Symmetric Attention, whose trainable graph matrices (shared weights over 2D grid graph edges) preserve chosen grid symmetries, applied locally (pixel-level windows) and globally (patch-level), with a directed-triangle sum that breaks flip symmetry while keeping 90° rotation invariance.
Key results
- On Procgen, SiT average normalized score 49.1% vs ViT 5.9% and CNN 30.4%; SieT reaches 51.7% average, comparable to UCB-DrAC ResNet with 9x parameters (48.2%).
- On CIFAR-10, GSA reaches 80% test accuracy vs 76% for global positional embedding and 56% without positional embedding, with SiT/SieT matching ViT performance using 2-5x fewer epochs.
- SiTs are 2x-5x slower to execute than ViTs of comparable size in the current implementation, due to custom layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use SiT instead of CNN-UCB | Procgen CaveFlyer · CaveFlyer | min-max normalized score: 18.0% -> 59.7% | Helped |
| use SieT instead of CNN-UCB | Procgen StarPilot · StarPilot | min-max normalized score: 44.6% -> 42.2% | Hurt |
| use SieT instead of CNN-UCB | Procgen Fruitbot · Fruitbot | min-max normalized score: 85.8% -> 76.0% | Hurt |
| use SieT instead of CNN-UCB | Procgen Chaser · Chaser | min-max normalized score: 44.6% -> 54.0% | Helped |
| combine invariant and equivariant representations in SieT | SiT on Procgen · Procgen average | min-max normalized score: 49.1% -> 51.7% | Helped |
| use two consecutive sums over triangles of the attention score matrix (ablation) | SiT on Procgen · Procgen average | min-max normalized score: 49.1% -> 50.6% | Helped |
| replace positional embedding with graph symmetric attention | ViT on CIFAR-10 · CIFAR-10 test | accuracy: 76% -> 80% | Helped |
| remove positional embedding (ablation) | ViT on CIFAR-10 · CIFAR-10 test | accuracy: 76% -> 56% | Hurt |
| add anti-symmetrisation to the attention score matrix (ablation) | SiT on MiniGrid grid-world · MiniGrid LavaCrossing | generalization performance | Helped |
| use the G and G_kq,b graph-matrix variant (ablation) | SiT on MiniGrid grid-world · MiniGrid LavaCrossing | performance | Mixed |
| include local GSA (ablation) | SiT on MiniGrid · MiniGrid | performance | Helped |
| use graph symmetric attention instead of positional embedding in the ViT | Transformer on CIFAR-10 · CIFAR-10 | training epochs to reach the same performance: 2-5× less training epochs | Helped |
Surprising findings
- Removing positional embedding entirely (full permutation invariance) degrades CIFAR-10 accuracy to about 56%, while GSA is better than positional embedding.
Limitations
- Local GSA adds substantial compute and memory cost, limiting scaling to larger models and complex tasks.
- Gains on Procgen are mainly in CaveFlyer, which has strong symmetry; other games lack symmetry, and single-task performance is comparable to CNNs.
- Atari and DM-control results are proof-of-concept with limited tuning and few games.
Experimental setup
Datasets: MiniGrid LavaCrossing Procgen Atari 100k DM-control Walker-walk CIFAR-10
Benchmarks: MiniGrid LavaCrossing Procgen (CaveFlyer, StarPilot, Fruitbot, Chaser, full 16-game suite) Atari 100k (5 games) DM-control Walker-walk CIFAR-10
Metrics: min-max normalized return test reward test accuracy sample efficiency
Baselines: CNN (ResNet) ViT E2CNN UCB-DrAC CNN (data-efficient, Rainbow) IRIS + CNN CNN (CIFAR-10 supervised ViT baseline)
Builds on: Vision Transformer rotary embedding graph attention networks DrAC IMPALA
Closest papers by meaning
- Boosting Sample Efficiency and Generalization in Multi-agent Reinforcement Learning via Equivariance (NeurIPS 2024 · similarity 0.87)
- Zero Shot Generalization of Vision-Based RL Without Data Augmentation (ICML 2025 · similarity 0.87)
- Focus On What Matters: Separated Models For Visual-Based RL Generalization (NeurIPS 2024 · similarity 0.87)
- A Practical Guide for Incorporating Symmetry in Diffusion Policy (NeurIPS 2025 · similarity 0.87)
- The Surprising Ineffectiveness of Pre-Trained Visual Representations for Model-Based Reinforcement Learning (NeurIPS 2024 · similarity 0.87)
- Foundation Policies with Hilbert Representations (ICML 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.