SiT: Symmetry-invariant Transformers for Generalisation in Reinforcement Learning

Matthias Weissenbacher, Rishabh Agarwal, Yoshinobu Kawahara
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:52695-52719, 2024.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-weissenbacher24a, title = {{S}i{T}: Symmetry-invariant Transformers for Generalisation in Reinforcement Learning}, author = {Weissenbacher, Matthias and Agarwal, Rishabh and Kawahara, Yoshinobu}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {52695--52719}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/weissenbacher24a/weissenbacher24a.pdf}, url = {https://proceedings.mlr.press/v235/weissenbacher24a.html}, 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.} }
Endnote
%0 Conference Paper %T SiT: Symmetry-invariant Transformers for Generalisation in Reinforcement Learning %A Matthias Weissenbacher %A Rishabh Agarwal %A Yoshinobu Kawahara %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-weissenbacher24a %I PMLR %P 52695--52719 %U https://proceedings.mlr.press/v235/weissenbacher24a.html %V 235 %X 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.
APA
Weissenbacher, M., Agarwal, R. & Kawahara, Y.. (2024). SiT: Symmetry-invariant Transformers for Generalisation in Reinforcement Learning. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:52695-52719 Available from https://proceedings.mlr.press/v235/weissenbacher24a.html.

Related Material