Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper presents a representation learning technique for data-efficient RL that makes use of invariance/equivariance of the environment wrt both state and agent actions. This work builds on two ideas: MDP homomorphism and symmetries in state-actions due to some inherent structure in the environment. MDP homomorphism refers to a mapping of states and actions of an MDP to a new space where the MDP reward structure and dynamics are preserved. Such a homomorphism can be useful if the new mapped state-action space is easier to work with. The second idea is that we may know that an environment exhibits certain regularities (i.e., may have a group structure) where we know the effect of certain operations. Then building these regularities directly into our models can then make training much more data efficient.

In this paper, the authors combine these two ideas (MDP homomorphism and symmetries in MDPs) such that state and actions are mapped to new spaces that satisfy the MDP homomorphism criteria along with the invariance/equivariance relations that result from the structure of the environment. This is achieved by minimizing a loss function that consists of 3 terms. The first term (L_AET) encourages the transition dynamics in the learned space to be the same with the original MDP. The second (L_GET) encourages the effect of group actions to be equivariant on the learned space. And finally the third one (L_R) encourages the rewards in the learned space to be the same as the original MDP.

The authors test their technique (in a Rainbow DQN agent) on Atari levels under data limited setting and show that their technique, albeit slighlty, outperforms previous techniques.
main_review: Overall, I think this is a quite interesting paper. It presents a very principled approach to learning better representations for RL in a data-efficient manner. Perhaps the improvement in performance is not very dramatic (and parts of this were already in earlier publications), but I think the motivation, the technical development, and the final presented technique are all of interest to the community and will be valuable contributions.

I don't have many concerns about the paper but I would have liked to see some analysis of the learned representations and the transition models. Are there any interesting patterns emerging? How does the new learned state/action space look like?

Also, it was not clear to me why adding the loss term L_GET led to worse results (in median scores). It would be nice to discuss this briefly. This might be perhaps the assumed group structure is not really suitable for some of the games. 

Some other points

- It would be nice to have colors in Figure 2. As it stands, it is a little bit hard to parse.
- It would be nice to mention what group SE(2) is at some point (perhaps in the Experiments section).
- L_AET is always included in training but this is only mentioned afterwards. It might be better to mention this earlier.


summary_of_the_review: Overall, I think the technique presented in this paper is well-motivated and addresses an important problem in RL. Perhaps the improvement in performance is small but I think the technical development and the final technique are valuable contributions.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper studies the question of recovering representations that are equivariant under the agent's action and under the symmetries of the environment. The representation are symmetry-based, using the state-transition symmetry (to create a strong inductive bias robust to symmetries that are not present in the reward) in a latent transition model (to enable learning of the transition model in the latent space, and to allow for linear assumption of the group action through the representation). A decomposed structure is imposed on the representation of the symmetry group, such that it allows to represent the latent embedding space as a direct product of subgroups. The authors propose 3 losses that encode the aforementioned constraints of the latent state and state-action embedding, as well as the invariance constraint of the reward.
main_review: I found the paper to be very well written. I think the idea and method were clearly explained despite the complexity of the overall model, combined with all the moving parts, well done! 

The author's perspective on equivariance in RL is interesting and novel (to the best of my knowledge). I was convinced by the motivations outlined and the consequently design choices. The proposed method makes sense, and is backed by interesting results.

I think it would benefit the paper to make it clear early on that the choice of the loss depends on a choice of subgroup blocks that depends on prior knowledge the environment (e.g. 2d translation, rotation, etc.) as this didn't come clearly to me until Example 1 is presented. I appreciate that the appendix includes a discussion on the choice of group for Atari games, but it would be better to have more general discussion in the main text. 

I understand that authors wanted to show performance across many different environments to showcase the robustness of the performance, but I would have liked to see more in-depth experimental results in one or two tasks including learning curves, specially given that data-efficiency is the main desired property. 


minor comment:
I believe the b_{\eta}  above Eq 18 should be g_{\eta}?
summary_of_the_review: The paper is well written and introduces a very interesting perspective on equivariant representations in RL. The method proposed makes sense to me, and the results seem promising.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper presents a latent variable model for representation learning in RL, taking into consideration both equivariance to an agent’s action and symmetry transformations of the environment. The proposed method (EqR) outperforms several previous methods in data-efficient (100k) Atari benchmark.
main_review: Strengths:

* Leveraging invariance and equivariance as inductive bias for representation learning in RL makes sense. The proposed way to encode the equivariance / invariance constraints in the latent representation space is novel and interesting.
* The writing is in general clear and not too difficult to follow.

Weaknesses:

* One of my concerns about this paper is the performance comparison with SPR. From Table 1, the performance of (EqR, $L_R$) is very close to SPR in the median score. Since on Atari median score is often considered as a more reasonable performance measure than mean score, I would not be convinced that EqR outperforms SPR.
* Besides, it seems (EqR, $L_R + L_{GET}$) performs worse (in terms of median score) than (EqR, $L_R$) in Table 1 and Figure 3(b). The novel part of this paper, $L_{GET}$, does not contribute much to the performance. Given that the method is also very similar to SPR in some implementation choices, I am not convinced that the proposed method is better.
* To help readers better understand the math in Section 3.2 and Section 4, I suggest the authors give more examples during these 2 sections to provide a context (possibly with the pendulum example).

Other comments:

* In Equation 5, can the same $g$ act both on element in $\mathbb{S}$ and on element in $\mathbb{S} \times \mathbb{A}$ ? 
* In the first paragraph of Section 4, $\mathcal{G}_T$ and $\mathcal{G}_R$ are not defined.
* I am not quite getting the sense of using $\tau_\theta(s_{t'})\tau_\theta(s_t)^{-1}$ as $\tau_g$. Can the authors provide some explanations?
* Typo: above Equation 17, duplicate "using".
* Above Equation 18, $b_\eta$ and $b$ are not consistent with Figure 2.
summary_of_the_review: While I do find the idea of this paper novel and interesting, the empirical evidence does not convince me of its superiority over prior methods, in particular SPR. I will not recommend acceptance for its current form.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper focuses on the problem of incorporating symmetries into RL agents’ representations to improve data efficiency. It proposes a model-based representation learning method parameterized by a group action that is based on encoding states and state-action pairs as elements of the group's representation in latent space. The proposed method is evaluated against a number of baselines in the data-limited Atari setting.
main_review: ### Strengths

- The paper is well-motivated and addresses an interesting and important problem.
- The idea of embedding states and actions as matrices representing a group action is intriguing, and may induce interesting properties in the representation beyond those studied in this paper.
- The paper provides sufficient background to be understood by readers with some but not extensive experience in the topic.


### Weaknesses

- **The justification for the proposed method is lacking** — section 4 provides a number of desiderata, but there is no follow-up showing empirically or theoretically that the resulting method (especially given the use of target networks and projections) will actually satisfy these desiderata.
- **The use of Lie groups is not motivated in Section 3**. As a reader, I presumed that this class was chosen in order to justify the use of learned matrix embeddings in the proposed model, but this would be worth clarifying. Additionally, the basis of the Lie algebra is not referred to in Section 6 despite an indication that it will be used in Section 3.

- **The paper suffers from a lack of clarity.** For example, in the discussion of symmetric MDPs the group action is applied to both states and state-action pairs, but it is not clear whether these represent two separate actions, or whether $g\cdot \langle s, a \rangle$ has the same effect on states as $g \cdot s$. Additionally, the notation $\kappa (\mathcal{G})$ and $\tau(\mathcal{G}_{\mathbb{S}})$ is not clear on a first read. In general I think more explicitly stating the types of the mathematical objects used in the paper will make it much easier to parse. By default, I and probably many other RL readers default to assuming vector-valued embeddings, and it is easier to override this default when the output type of the embedding is made explicit.
	- Section 4 is quite long and in many cases it’s not clear whether the proposed properties are _desiderata_ or guaranteed to be satisfied by the proposed method. For example, the matrix embeddings of state actions seems to hold by construction but disentanglement presumably does not.
- **The empirical evaluations don’t support the main goal of the paper.** Adding the group-equivariant loss doesn’t provide much improvement over the standard model-learning loss — particularly when we look at median performance. This suggests that most of the benefit over the baselines is coming from the model-learning component of the objective rather than the incorporation of symmetries. Further, the method without the symmetry-learning loss is similar to many existing representation learning methods such as DeepMDP, PSR, PBL, MuZero, LatCo, and DBC. The equivariant loss is therefore critical to the novelty of the proposed algorithm, yet it doesn’t seem to have a significant impact on overall performance. It is possible that a more careful empirical analysis on environments with more explicit symmetries may reveal that this loss is indeed useful, but the current experimental results are insufficient to show this to a satisfactory degree.
    - The experimental results are limited. The paper only uses one evaluation scheme and doesn’t provide much insight into whether the learned model really is equivariant to the set of transformations used. 
    - Further, the baselines used in the paper don’t seem to be described or cited, which makes it difficult to determine whether a fair comparison is really being made between the different methods. The baselines I was already familiar with mostly seem to make use of model-free approaches combined with data augmentation, which leaves the model-learning component of EqR as a confounding factor. SPR, which does include a learned latent-space model, seems to attain a higher median human-normalized score than EqR with the equivariant transition loss, suggesting that at least part of the improved performance is coming from the specific parameterization of the model. 

### Concrete avenues for improvement

- I am very intrigued by the use of matrix embeddings in the representation — typically in RL we treat representations as vectors, and it would be interesting to see whether enforcing linear actions of the transition and symmetry operators on the latent space induces particular structure. 
- It is possible that by looking at environments with more structure, the proposed method might be shown to capture the appropriate in/equivariances and therefore improve performance and sample efficiency. I recommend incorporating evaluations on e.g. Mujoco or some other environment with more explicit physically-based invariances. 
- The focus of this paper is on equivariant representations for sample efficiency; however, it seems like an arguably more promising application is in improving generalization. 
- Ensuring that symbols and their types are defined before they are used will make the paper much more readable. 
- Section 4 should be rewritten to more clearly express which of the desiderata are definitely satisfied by the proposed method, and which are not guaranteed to be satisfied. It might also be interesting to include empirical investigations into the degree to which the proposed method satisfies these desiderata. 
summary_of_the_review: I recommend rejecting this paper due to a) the limited empirical benefit of incorporating equivariance into the proposed representation learning method and b) the lack of clarity in both the paper’s exposition and its claims about the proposed method’s properties. With additional work and a more informative evaluation setting, I think the paper will be in a much better position for consideration in future conferences.
review_title: null

Meta-review: This paper proposes to learn a latent space representation such that some linear equivariance and symmetry constraints are respected in the latent space, with the goal to improve sample efficiency. One core idea is that the latent space is also the same as the space of linear transformation used in the constraints, which is shown to simplify some of the mathematical derivations. Experiments on the Atari 100K benchmark demonstrate a statistical improvement over the SPR baseline when using the SE(2) group of linear transformations as latent space.

Following the discussion period, most reviewers were in favor of acceptance. However, one reviewer remained unconvinced, and after carefully reading the paper, I actually share the same concerns, i.e., that it is unclear under which conditions the proposed approach actually works, and what makes it work. I believe that, as a research community, we should value understanding over moving the needle on benchmarks, especially when proposing such a complex method as this one (see Fig. 5).

More specifically:

1. The method is only evaluated on Atari games, showing some improvements when using SE(2), and arguing that there are corresponding symmetries in such games. There is however no analysis demonstrating (or even hinting at the fact) that the proposed technique is actually learning to take advantage of such symmetries (NB: I had a quick look at the animation added by the authors in the supplementary material, but I do not see if/how they help on this point). Even if analyzing representations on Atari may be tricky, I believe that given the motivation of this new algorithm, it *must* be evaluated on some toy example (e.g., the pendulum mentioned throughout the paper) to validate that it is learning what we want it to learn (although I also agree with the authors that experimenting on a more complex benchmark like Atari is equally important).

2. The idea of embedding states into the same space as transformations is interesting, and brings some advantages when writing down equations, as demonstrated by the authors. However, there is no justification besides mathematical convenience, and it doesn't seem intuitive to me at all that why this should be a good idea, considering that it ties the state representation to the mathematical representation of group transformations. For instance, what does the spcial group element $e$ mean for a state? And this coupling makes it difficult to interpret the effect of using a different group of transformations: for instance when moving from GL(2) to SE(2), is the observed benefit because we are using only specific transformations, or simply because we are reducing the dimensionality of the state embedding? (note that in Fig. 4(c) the MLP variant has similar performance to GL(2), and based on my understanding they use the same embedding dimensionality  ==> I believe it would be important to check what would happen with an MLP variant using the same dimensionality as SE(2))

3. The effect of the $L_{GET}$ loss is not convincing, as pointed out by several reviewers. I think it would have been an opportunity for the authors to investigate why, especially since it seems to work in some games and not others. But just focusing on "here are the 17/26 games where it works better" doesn't really bring added value here. Do these games have some specific properties that make them better candidates to take advantage of $L_{GET}$? This could have been a very interesting insight if that was the case, but as it is now, I am not sure what we can learn from that.

4. There are several implementation "details", some moving the final algorithm farther from its theoretical justification, that are not ablated, making it difficult to understand their impact (ex: using target networks, the choice of the value of M, using projections onto the unit sphere of some arbitrary dimensionality, how the $s'$ state is chosen in $L_{GET}$)

As a result, we have here an algorithm with some interesting theoretical background, but with a lot of moving components which -- when properly tweaked -- can lead to a statistically meaningful improvement on Atari 100K -- without really understanding why. I believe this is not quite enough for publication at ICLR, and I would encourage the authors to delve deeper into the understanding of their algorithm, which I hope will bring useful insights to the research community working on representation learning.
Decision: Reject