SiT: Symmetry-invariant Transformers for Generalisation in Reinforcement Learning

...inevitably hurts sample efficiency and usability as compared to other commonly used vision encoders in vision-based RL tasks. No learning curves or results are provided related to the sample efficiency for RL tasks, which is somewhat insufficient to demonstrate the practical value of the proposed method to the RL community.

• Indeed, we agree with the reviewer that sample efficiency curves provide a lot of insight. All SiT variants demonstrate comparable sample efficiency to the CNN baseline, see Figure 7a. We have incorporated the training and testing curves for the main Procgen experiments into the rebuttal version. Transformers have enabled great advances in nlp and vision, particularly for scaling and our work possibly opens these avenues for image-based RL.

symmetry-preserving designs in CNN or other image-based models. Such methods should be also compared

• We have added the E2CNN baseline and find SiTs outperform this baseline. We emphasize that aside from our performance gain we are the first to succeed in training ViTs (SiTs) in RL on Procgen.

How will the proposed method perform on more general visual control tasks, like Atari?

• Atari doesn't allow for studying generalization easily, which is the focus of this work while procgen is widely used given its wide variety of levels, and a more complex action-space, see e.g. [muZero].

---

---

1. ...the paper should be treated as a CV paper rather than an RL paper…

• Relevance of Model Evaluation in RL: Many studies in Reinforcement Learning (RL) focus on model architecture evaluation, including works directly relevant to ours such as those by Wang et al., and Ha et al. Other reviewers have also mentioned more examples in this context.

2. ...For some tasks like robotic manipulation, as… the importance of symmetry-invariant may not be that large…

• Actually, invariant and equivariant networks [Wangx2] have been shown to give benefits in robotic control.

3. ...ViT is quite heavy and costly to learn, it inevitably hurts sample efficiency and usability as compared to other commonly used vision encoders in vision-based RL tasks…

• (a) Transformers and ViTs in Vision-Based RL: We agree, Transformers and Vision Transformers (ViTs) are more computationally intensive and less sample efficient, which has limited their application in vision-based RL. This challenge was a primary motivation for our research. We believe our approach shows promise in altering this landscape.

• (b) Given the dominance of ViTs in Vision and Transformers in NLP, it's plausible that improvements in Transformer technology will similarly revolutionize vision-based RL, with ViTs becoming predominant. Recent technical advancements, such as efficient Transformers [flashattn], which offer up to a 10x performance boost may lead the way.

…The proposed method is based on ViT, which makes its applicability only restricted to vision-based RL tasks…

• (a) Vision-based RL is a pivotal subfield of deep RL, marking the origin and significant advancements in the field. The focus of our paper on this area upholds the relevance and depth of our research.

• (b) Other Application of Transformers in RL: Moreover, transformers are central to Decision Transformers [a] and the recent IRIS [b]. Our approach is also applicable to these models. We originally chose not to elaborate further on 1D-data to not overload the paper conceptually. However, it would be a minor change to add a paragraph on 1D-data of SiTs. see appendix A.1 where we have added such a brief discussion.

---

---

Minor: the paper reminds me of a recent NeurIPS paper [1] that also uses symmetry to enhance RL performance.

• We have added the suggested reference, even though we concur that it is somewhat orthogonal to our paper.

---

References

[muzero] Procedural Generalization by Planning with Self-Supervised World Models, A.Anand, J. Walker, Yazhe Li, Eszter Vértes, Julian Schrittwieser, Sherjil Ozair, Théophane Weber, Jessica B. Hamrick :https://arxiv.org/abs/2111.01587

[Wangx2] Dian Wang and Robin Walters. So (2) equivariant reinforcement learning. In International Conference on Learning Representations, 2022 Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics. PMLR, 2022b.

[flashattn] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré https://arxiv.org/abs/2205.14135

...The authors should use an E2 equivariant architecture like E2 Steerable Networks as their baseline for RL experiments…

• We have added the E2CNN baseline, which outperforms CNNs on Caveflyer and Chaser, and find that SiTs outperform this stronger baseline. We emphasize that aside from our performance gain, we are the first to succeed in training transformer-based models ViTs (SiTs) on Procgen RL tasks. in RL on Procgen. Transformers have enabled great advances in NLP and vision, particularly for scaling and our work possibly opens these avenues for image-based RL.

Some sections of this paper are not well written leading to confusion … for example in equation 2, symmetric(GV) assumes that G is $\mathbb{R}^{P \times P}$ whereas according to the author's definition of G it should be $\mathbb{R}^{P \times d_f}$.

• We have made the background section more concise in the rebuttal version, in particular the discussion around equation 2, which may have caused this confusion. We define G as $G \in \mathbb{R}^{P \times P}$. See e.g., in the definition of G page 4, or Figure 2, which visualizes G . Also we could not locate any such typo in our manuscript. We do apply the graph matrix per each feature so there is an additional $d_f$ in the dimensionality as $G \in \mathbb{R}^{P \times P \times d_f}$.

...Authors should also expand on the related work on equivariant architectures [5,6] for reinforcement learning. [3, 4, 7]

• We have included the suggested references, and mentioned them in the related work section. While E2CNN networks abide by equivariance they are based on CNNs in contrast to our ViT based approach, which moreover allows for invariance, equivariance as well as a mixture thereof i.e. SieT.

---

Answers to Questions

In Figure 1 (a), what does the bottom right image depict?

• Figure 1 (a) bottom-right shows the original image where a minimal number of patches are permuted (six to be precise). This highlights that even minimal permutation invariance of an agent may be fatal for learning - aligned with our motivation that PI generally should be broken by some mechanism i.e. GSA. We have clarified this in the revision

Can you explain how Equation 2 is connected to GAT and E2GNN?

• Equation 2 is given to provide an overview formulation of GAT. There is no direct connection of Equation 2 to E2GNN (Satorras et al., 2021b). We have clarified this in the revision.

...This makes it really hard to follow the author's argument and how this formulation is connected to Graph Attention Network…

• It appears the above misunderstanding may have been caused by the appearance of “G” in equation 2, which is not the G introduced in our work e.g. in Figure 2. We have changed the notation to $\mathcal{G}$ in equation 2 for the adjacency matrix.

We thank the reviewer for taking the time to provide their helpful and valuable feedback. Our response follows:

The experiments could benefit from stronger baselines. …. like e2cnn.

• We have added the E2CNN baseline, which outperforms CNNs on Cavelfyer and Chaser, and find that SiTs outperform this stronger baseline. We emphasize that aside from our performance gain, we are the first to succeed in training transformer-based models (SiTs) on Procgen RL tasks. Transformers have enabled great advances in NLP and vision, particularly for scaling and our work possibly opens these avenues for image-based RL.

it would be nice to have an ablation study on only using GSA globally or locally…

• On Minigrid Lavacrossing, SiT with 2 global GSA layers falls short compared to 1 local & 1 global and 1 local & 2 global layers, respectively; in particular, on the most challenging generalization test tasks i.e. random goal and hard (N3) environment. For reference, please see the appendix section G1, particularly Figure 10 and Table 4, where we have explored this aspect. However, isolating the effects of local GSA alone presents architectural design challenges. But the local effects can be inferred indirectly by comparing 1 loc./2 glob. to 2 glob. in Figure 10 and Table 4.

---

Answers to Questions

... I think G should be an all-one matrix here instead of an identity matrix…

• In order to get the conventional attention mechanism, G needs to be the identity matrix. Please review Equations 17 and 18 in the appendix, where we detail the index contractions of GSA. Since we apply graph matrices for each feature, the identity matrix is replicated for the number of features $d_f$. The intriguing observation that this corresponds to a graph without connections to other vertices aligns nicely with our graph interpretation of G.

I do not fully understand why $G_{v,k,q}$ are necessary…. Would just having G in equation 3 not be enough for maintaining the desired equivariance?

• Correct, both $G$ and $G_{v,k,q}$ are indeed sufficient to break symmetries, respectively, i.e. to achieve the desired equivariance. We delve into this more in the paragraph titled Explicit & Adaptive Symmetry Breaking on page 5. In our preliminary experiments, G required more h-param tuning- probably due to the Hadamard product -, thus we employ $G_{v,k,q}$. We have added this discussion to the paper.

How are the sizes of the local patch (5x5) and the global patch (14x14) selected?

• For Minigrid, we selected a 14x14 global image size as this is the minimum resolution at which the direction of the triangular shaped agent can be discerned when downscaling the RGB rendering of the environment. Smaller resolutions fail to capture this detail, while larger ones are feasible but not necessary. The local neighborhood size in Minigrid (in pixels) is required to be odd, so 3x3 and 5x5 are the smallest viable options. We have revised the paper to include this detail.

We're grateful for your acknowledgement of our improvements and thank you for the score increase!

For example, papers like [1,2] should be added in related work and it should be discussed how the current technique in comparison to those and argument about why they did not compare with those.

We thank the reviewer for their suggestions. To address their main concern, we have updated the related work section to better position our work with respect to prior work on symmetry in RL.

While algebraic symmetries in RL and symmetry-based representation learning can improve sample efficiency in a single environment using CNN-based architectures as backbone; the focus of our work is on improving generalization in RL by incorporating symmetries in vision transformers. Due to the widespread use of transformers in NLP and vision, transformer-based architectures carry many additional promises for future development in RL, see our revised conclusion section for examples. Interestingly, SiTs are complimentary and may be added on top of some prior works incorporating symmetries e.g. [1,2], which we leave for future work.

if the authors have Atari 100k experiments they can compare with [2] and show how their method performs.

Based on the discussion with reviewer Rs9f, we ran some Atari 100k experiments to serve as a proof-of-concept that SiTs can be extended easily to other tasks. That said, we want to reiterate that the focus of our work is on generalization in RL as opposed to data efficiency.

Specifically, we train SiT on top of Rainbow without h-parm changes, and find that SiTs compare similarly to CNNS, see Table in section F.3. To the best of our knowledge, this is the first time ViTs + Q-learning has been successfully applied for online RL on Atari.