Paper: SiT:   Symmetry-invariant Transformers for Generalisation in Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The majority of the paper is to derive a symmetry-invariant and equivariant ViT model, and there is not much design in the RL part. That's why I think the paper should be treated as a CV paper rather than an RL paper. In that sense, the proposed method should at least first demonstrate its superiority in CV tasks. Unfortunately, the proposed method is only evaluated on the extremely simple CIFAR-10 dataset, and compared with no other CV baselines except for ViT.

Author response:
**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.

**CV focus:** We would be grateful if you could further elucidate your viewpoint, particularly how our contribution contrasts to other RL paper doing model evaluation  (several cited in our work & published as RL paper in major conferences )?

Reviewer comment (Evaluation):
The proposed method is based on ViT, which makes its applicability only restricted to vision-based RL tasks. Moreover, as 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. 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.

Author response:
**...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. 


**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.**

Reviewer comment (Experiments):
The evaluations are only conducted in tasks that particularly rely on image generalization capability, like Procgen and MiniGrid. For some tasks like robotic manipulation, as the sense of orientation is vital, the importance of symmetry-invariant may not be that large.

Author response:
**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.

Reviewer comment (Experiments):
There lack of strong baselines in the experiments. The baselines only include ViT, variants of the proposed method, and simple baselines like CNN. As far as I know, there are also some existing works that introduce symmetry-preserving designs in CNN or other image-based models. Such methods should be also compared given their relevance to this paper.

Author response:
**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.

Reviewer comment (Novelty):
Minor: the paper reminds me of a recent NeurIPS paper [1] that also uses symmetry to enhance RL performance. This is probably not super related as it enforces symmetry on the temporal aspect, but also worth mentioning.

Author response:
**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.

Reviewer comment (Experiments):
How will the proposed method perform on more general visual control tasks, like Atari?

Author response:
**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].  

**Atari:**  We appreciate your suggestion to include Atari tasks for a more comprehensive evaluation.   Could you please elaborate on how Atari tasks might provide a different measure of generality compared to Procgen? Procgen and MiniGrid as they are widely recognized benchmarks for testing generalization in vision-based RL tasks. Note that Fruitbot and Starpilot do not rely on symmetries, so they are not specifically chosen but probably the most discussed Procgen environments. Note that we do have a proof of concept on DM-control walker-walk in the paper. ) 

**We have incorporated four Atari evaluations as per your request. We hope this addition addresses your concerns effectively. Thank you for advocating for this inclusion; we believe it significantly enhances the value of our paper.**

Reviewer comment (Experiments):
The experiments could benefit from stronger baselines. Given the paper's introduction of a novel equivariant architecture, I think a comparison with existing equivariant architectures is necessary. A possible baseline could be some equivariant architectures that enforce global rotational symmetries like e2cnn [A], or utilizing rotation data augmentation to realize equivariance in CNN or ViT.

Author response:
**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.

Reviewer comment (Experiments):
It would be nice to have an ablation study on only using GSA globally or locally to understand how the two components contribute to the improvement.

Author response:
**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.

Reviewer comment (Theory):
In the first paragraph of page 4, the paper claims, `Notably, GSA reduces to the conventional attention mechanism when the underlying graph topology matrix G only contains self-loops, i.e. G being the identity matrix.` My interpretation of this is that if $G$ is the identity matrix, the output of $\Gamma(Q, K)$ would result in a diagonal matrix, where each token can only attend to itself. This doesn't align with the conventional attention mechanism. I think $G$ should be an all-one matrix here instead of an identity matrix. Please correct me if I am missing something here.

Author response:
**... 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.

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

Author response:
**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.

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

Author response:
**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.

Reviewer comment (Writing):
Some sections of this paper are not well written leading to confusion while following the paper's arguments. For example in equation 2, symmetric($GV$) assumes that GV is $\mathbb{R}^{P\times P}$ whereas according to the author's definition of $G$ it should be $\mathbb{R}^{P\times d_f}$. This makes it really hard to follow the author's argument and how this formulation is connected to Graph Attention Networks [1] or $E(n)$ Equivariant GNN [2]. This also makes following section 3 difficult.

Author response:
**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}$. 

**...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.

Reviewer comment (Experiments):
As SiT has been built keeping in mind the inductive biases coming from the symmetries of the environment and task, I think just comparing with CNN or ViT baseline is not completely fair. The authors should use an E2 equivariant architecture like E2 Steerable Networks as their baseline for RL experiments [3, 4] or E(n) equivariant GNN [2]. Authors should also expand on the related work on equivariant architectures [5,6] for reinforcement learning. [3, 4, 7]

Author response:
**...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.

**...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.

Reviewer comment (Presentation):
In Figure 1 (a), what does the bottom right image depict?

Author response:
**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

Reviewer comment (Theory):
Can you explain how Equation 2 is connected to GAT and E2GNN?

Author response:
**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.