Paper: Can Euclidean Symmetry Help in Reinforcement Learning and Planning? (ICLR-2024)

Reviewer comment (Theory):
The paper does not adequately discuss two important aspects of equivariance in robotics: What if the equivariance is latent, i.e. the observations are not given in known, simple representations of the symmetry group – instead, they are given in pixels? And what if the symmetry group is partially broken, for instance by object positions or the direction of gravity?

Author response:
Your question about latent equivariance and partial symmetry breaking in robotics is indeed a fascinating area of research. In our study, we focused on scenarios where symmetry is known and explicit, as this forms the foundational basis for our theoretical exploration. While integrating latent equivariance or dealing with partially broken symmetry groups, such as those influenced by object positions or gravity, is highly relevant, we considered it beyond the scope of our current research. This is partly because (1) understanding the usefulness of symmetry in known cases is a preliminary step before delving into more complex scenarios, and (2) the problem of encoding unknown / latent / partially broken symmetries has been addressed in other works, such as Park et al., ICML 2022. However, we acknowledge that this is an important direction for future research and intend to explore it in subsequent studies.

Reviewer comment (Novelty):
The authors miss some references on equivariance in model-based RL and planning:  
	- A. Deac at al, "Equivariant MuZero", arXiv:2302.04798  
	- J. Brehmer et al, "EDGI: Equivariant Diffusion for Planning with Embodied Agents", arXiv:2303.12410

Author response:
- The two mentioned papers are indeed relevant but also quite new. We will discuss them in the related work.

Reviewer comment (Experiments):
In Section 4, the proposed method require a "G-steerable equivariant MLP" for continuous groups G. What architectures do you have in mind? If I'm not mistaken, you never experiment with any such architecture for continuous G, right?

Author response:
Linearization of G-steerable equivariance constraints requires continuous G. We use the theory to show that it’s useful to have continuous group in continuous action space. The sampling-based algorithm is designed for continuous action, and can use discrete groups to still benefit sampling-based algorithm.

In short, continuous group + continuous action is required to derive theory, while we can use discrete group for sampling-based algorithm for continuous action space.

Reviewer comment (Theory):
To make MPC equivariant, the action sampling needs to be equivariant. Just after Eq. (11), you describe as the main problem that "action sampling is not state-dependent". Why? Isn't the equivariant learned policy used for action sampling?

Author response:
In our framework, there are two main components to consider in the context of action sampling:

1. **Policy Network:**
    - The policy network plays a significant role in our approach, contributing approximately 25 samples in each Cross-Entropy Method (CEM) iteration (based on 512 total samples with about 5% coming from the policy network). This policy network is crucial for guiding the action sampling process towards promising regions of the action space.
2. **Sampling from a Fitted Gaussian Distribution:**
    - In addition to the policy network, we also sample actions from a Gaussian distribution fitted to the top-performing samples. Each CEM iteration involves sampling 512 actions from this distribution. The purpose of CEM here is to iteratively adjust the Gaussian distribution (mean and variance) towards an optimal set of actions for the given state.

However, a key point to note is the initialization of the Gaussian distribution. For each sampling process, we initialize the mean and variance from a standard normal distribution, N(0,1)*N*(0,1), and iterate this process six times. This initialization approach means that the mean and variance of the Gaussian distribution are not directly dependent on the current state of the system.

The rationale behind this choice is to maintain a balance between exploration and exploitation. By not making the Gaussian parameters directly state-dependent, we allow for a degree of exploration which is critical in scenarios where the optimal action is not immediately apparent from the current state alone.

Reviewer comment (Writing):
I wasn't able to follow the logic behind the G-sample method, could you explain that a bit more slowly, please?

Author response:
1. **Equivariance in the G-sample Method:**
    - In our Model Predictive Path Integral (MPPI) framework, we sample $N$ actions from a Gaussian distribution  $N(\mu, \sigma^2I)$, creating a set $A = \{a_i\}^N_{i=1}$ . A key aspect of our method is ensuring that the action sampling process maintains equivariance.
    - Equivariance necessitates that any transformation of the input state (for instance, a rotation) should correspond to a similar transformation in the output. However, since the parameters $\mu$ and $\sigma$ of the Gaussian distribution in standard MPPI are not state-dependent, they do not naturally satisfy this condition.
2. **Augmenting Action Sampling for Equivariance:**
    - To address this, we augment the action sampling by transforming it with all elements in the group $G$, forming a new set $GA = \{g \cdot a_i | g \in G\}^N_{i=1}$.
    - This augmentation ensures that our method can handle various orientations, maintaining the property of equivariance. For each action $a_i$ in the original set $A$, the augmented set $GA$ includes the transformed action $g \cdot a_i$ for every element $g$ in $G$.
3. **Invariance and Closure Under Group $G$:**
    - The invariance property we aim to achieve is that the return of a trajectory, $\text{return}(\tau)$, remains unchanged under transformations in $G$. This means the trajectory's return is invariant to rotations or other symmetries defined by $G$.
    - By augmenting the action set $A$ with transformations from G, we ensure that the set of potential actions (the sample space) is closed under $G$. This closure means that applying any transformation from $G$ to an action in our set results in another action that is also in our augmented set $GA$.

This methodical approach is crucial for maintaining consistency across different transformations and is a key aspect of achieving both equivariance and invariance in our planning algorithm.

We have detailed the mathematical foundations and further explanation of this process in Appendix E of our paper. We hope this comprehensive response clarifies the G-sample method's logic, implementation, and its critical role in our research.

Reviewer comment (Experiments):
Is it fair to compare non-equivariant and equivariant methods with equal parameter count? Usually, equi methods have less parameters than non-equivariant counterparts with similar capabilities. This choice thus strikes me as a bit unfair to the baselines.

Author response:
1. **Equal Parameter Count in Our Comparative Analysis:**
    - In our study, we have intentionally used an equal number of free parameters for both equivariant and non-equivariant methods. The primary difference lies in the structures and inductive biases of the architectures rather than the quantity of parameters. This approach follows a widely accepted convention in the field for comparing different equivariant neural networks, allowing for a more direct assessment of the impact of the architectural differences on performance.
2. **Equivariance in Common Architectures:**
    - We would also like to point out that while certain architectures are typically referred to as non-equivariant, many common architectures inherently possess equivariant properties. A notable example is the Convolutional Neural Network (CNN), which exhibits translation equivariance. Thus, the distinction between equivariant and non-equivariant methods can be more nuanced than it initially appears, as many standard architectures integrate forms of equivariance in their design.
3. **Focus on Architectural Differences:**
    - Our comparison aims to highlight how the specific structural features and inductive biases of equivariant architectures contribute to their performance, especially in contexts where symmetry properties play a crucial role. By maintaining an equal parameter count, we strive to ensure that any observed performance differences are attributable more to these architectural aspects than to the sheer number of parameters.

We believe this approach provides valuable insights into the efficacy of equivariant architectures and contributes to a more nuanced understanding of their role in various applications.

Reviewer comment (Experiments):
The proposed algorithm is evaluated on only four environments, two of which (2D point mass and 3D point mass) are toy problems. In fact, the 2D point mass is mainly used for debugging purposes and is rarely reported in a scientific paper due to its simplicity. I suggest the authors incorporate more experiments either from the robotics literature or works on continuous symmetry in RL, such as [3].

Author response:
- **Choosing baselines and tasks.**
    
    In response to your suggestions regarding the expansion of our experimental setup, we are considering incorporating additional tasks from reference [3] into our study. This inclusion will allow us to demonstrate the applicability and efficacy of our approach in a broader range of environments, particularly those that emphasize continuous symmetry in reinforcement learning.
    
    Regarding the inclusion of additional baselines, specifically Dreamer as cited in [2], we wish to clarify that our primary focus lies in demonstrating the usefulness of equivariance in model-based reinforcement learning, particularly in continuous control contexts. The approach we are based on, Temporal Difference Model Predictive Control (TDMPC), has been shown to outperform Dreamer in continuous control tasks. While we acknowledge the importance of comparing various model-based RL baselines, our objective is to highlight the specific benefits that equivariance brings to these algorithms, rather than to conduct a broad comparison across different model-based RL methods. This focused approach enables us to more effectively illustrate the unique advantages that our equivariant framework offers in handling continuous control challenges.

Reviewer comment (Novelty):
The authors have overlooked some essential papers, and I have identified some of their stated contributions as overstated. On page 2 (first paragraph), they state that their approach expands on earlier research on planning on 2D grids, yet [4] has already examined the equivariant Muzero which integrates symmetries in a complex model-based planning algorithm. Furthermore, they assert that they are extending equivariant model-free RL to continuous states and actions, while [1] has already accomplished this in a broader context, not solely restricted to Euclidean symmetry. Some examples of other missing references are [5, 6].

Author response:
- **Connection to equivariant MuZero.**
    
    Regarding, equivariant MuZero, it is not just about “complex” model-based planning algorithm. MuZero is based on MCTS for action selection. We focus on continuous control and continuous symmetry.

Reviewer comment (Novelty):
The definition of Geometric MDP appears to be a rebranding of MDP homomorphisms [7], which was also extended to continuous states and actions [1]. It is not clear why the authors have chosen to rename a well-studied concept in the literature by adding some restrictions on the group symmetry. This can be very misleading to an inexperienced reader.

Author response:
- **Distinction of Geometric MDP:**
    
    The uniqueness of Geometric MDP in our work lies in the specialized application of MDP homomorphism principles, tailored to leverage the inherent symmetries in model-based reinforcement learning. Unlike traditional MDP homomorphism works, which often lack practical algorithms for exploiting symmetries, our Geometric MDP framework is designed with practicality and applicability in mind. This approach allows us to transcend the theoretical limitations of traditional MDP homomorphisms and offer concrete, actionable methodologies.
    
    Key aspects that distinguish Geometric MDPs in our work include:
    
    1. **Practical Application of Symmetries:**
        - Our focus is on the practical use of symmetries within MDPs. By narrowing our scope to specific, symmetry-based problems, we provide more practical algorithms that do not rely on NP-hard orbit searches or known dynamics. This practical orientation is essential for making the concept of MDP homomorphism more accessible and applicable in real-world reinforcement learning scenarios.
    2. **Equivariance Constraints:**
        - A cornerstone of our approach is the use of equivariance to impose constraints on solutions, facilitating "value equivalence" in a manner that is both theoretically sound and practically viable. This approach allows us to exploit symmetries without the need for exhaustive searches or a priori knowledge of the dynamics.
    3. **Planning in Quotient MDP M/G:**
        - Our methodology involves planning in a quotient MDP M/G, which is feasible only when there is a continuous group action on the MDP state/action space. This quotient map, discussed in detail in Appendix, is a unique feature of our approach that leverages continuous group actions to provide a more structured and symmetrical framework for decision-making.
    4. **Focus on Continuous Symmetry:**
        - In contrast to prior works that primarily focus on discrete symmetry, our research introduces and develops a theory for continuous symmetry. This theory is pivotal for algorithms designed for sampling-based methods in continuous action spaces, representing a significant advancement in the field.
    
    In summary, Geometric MDPs in our work represent a novel approach to MDP homomorphism, tailored to the demands and challenges of modern reinforcement learning. By bridging the gap between theoretical concepts and practical application, we offer a framework that is both innovative and applicable, addressing the complexities of continuous control tasks and the need for efficient, symmetry-aware decision-making algorithms.

Reviewer comment (Novelty):
The contributions of this paper appear to be incremental with respect to the prior work of [7] which explored the use of symmetry in model-based planning.

Author response:
- **Relation to MDP Homomorphism, Choice of Baselines, and Incremental Contributions:**
    - **Clarification on MDP Homomorphism:**
        - It appears there may be a misunderstanding regarding our work's connection to MDP homomorphism. While our research is indeed based on the concept of MDP homomorphism, we emphasize that the original MDP homomorphism works do not provide practical algorithms for leveraging symmetries. Our approach aims to expand the practical application of MDP homomorphism by focusing on a smaller, yet significant and practical subset of problems where MDP homomorphism or isomorphism is induced by symmetry in the MDP. This focus allows us to propose more practical algorithms without the need for NP-hard orbit searches and known dynamics, which are significant limitations in the traditional application of MDP homomorphism. In our work, we use equivariance to constrain solutions within the framework of homomorphism, achieving "value equivalence" in a more practical and applicable manner.
    - **Quotient MDP M/G and Continuous Group Action:**
        - As discussed in Appendix B.1, our methodology involves planning in a quotient MDP M/G, made possible only when G has a continuous group action on the MDP state/action space. This quotient map is central to our approach and differentiates our work from prior studies.
    - **Differentiation from Prior Work [8] and Expansion of Discussion:**
        - We acknowledge that our contributions extend beyond the work presented in [8], which focused on discrete symmetry. Our research introduces a theory for continuous symmetry and an algorithm tailored for sampling-based algorithms in continuous action spaces. We have expanded the discussion of symmetry in MDPs and MDP homomorphism in a way that provides both theoretical depth and practical application, addressing a gap in the existing literature.

Reviewer comment (Theory):
One of the key contributions of the paper, as claimed by the authors, is the study of continuous group symmetries in RL. Unfortunately, in their practical algorithm they are using discretized subgroups (page 8). This raises doubts regarding the soundness of the paper and the connection between theoretical analysis and the experimental results.

Author response:
- **On using discrete symmetry.**
    
    In short, continuous group + continuous action is required to derive theory, while we can use discrete group for sampling-based algorithm for continuous action space.
    
    More concretely, linearization of G-steerable equivariance constraints requires continuous G. We use the theory to show that it’s useful to have continuous group in continuous action space. The sampling-based algorithm is designed for continuous action, and can use discrete groups to still benefit sampling-based algorithm.

Reviewer comment (Theory):
The MDP dynamics is assumed to be deterministic (page 5, second paragraph) without any justification or insight into its reason.

Author response:
- **Stochastic dynamics.**
    
    The sampling-based algorithms do not need determinisiticity. Only the theory assumes determinisiticity, while it can be easily extended to Linear Quadratic Gaussian and considers stochasticity.

Reviewer comment (Theory):
The proposed method assumes the symmetry group is known, which may limit its practical application.

Author response:
- Our study primarily concentrates on equivariance between reference frames in D-dimensional spaces, focusing on distance-preserving transformations. Normally, the control tasks’ dynamics we are interested in are symmetric under these  distance-preserving transformations. The cornerstone of this approach is SO(D) equivariance, which we believe is crucial for understanding and leveraging symmetries in these contexts.

Reviewer comment (Experiments):
Since the proposed method extends previous work from discrete case to continuous case, can the proposed method also cover the tasks with discrete actions?

Author response:
- Technically, our approach can be extended to tasks with discrete actions. However, our current focus is on continuous control tasks where we utilize sampling-based algorithms like Cross-Entropy Method (CEM) and Model Predictive Path Integral (MPPI). These algorithms iteratively update a Gaussian distribution, which is well-suited for our study. We acknowledge that other parameterizations or methods like Monte Carlo Tree Search (MCTS) could be used, and this is an area we might explore in future work.

Reviewer comment (Experiments):
How is the performance of the proposed method compared to other well-known RL algorithms, e.g., SAC, and DDPG on DeepMind Control suite?

Author response:
- We understand the importance of a comprehensive performance comparison with other well-known RL algorithms. It’s worth to note that TDMPC already outperformed SAC in their original paper. We also included a “planning-free” equivariant TDMPC to mimic equivariant SAC, which was outperformed by the version with planning, as shown in Figure 11, “Ablation on model-based vs. model-free (“planning-free”)”.

Reviewer comment (Presentation):
The paper lacks focus on multiple fronts, as various Theorems, Propositions, paragraph titles, and types of lists make up the main body but do not give it structure, since the parts do not naturally follow from one another.

Author response:
Your suggestion is helpful for our reorganization of the paper. In response to your feedback, we are actively considering reorganizing the paper to better motivate the paper with our empirical study.

Reviewer comment (Evaluation):
The empirical study motivates further research but lacks a definite conclusion, and it is unclear how the described behavior translates to other means of training for MDPs.

Author response:
We are enhancing our empirical study to provide clearer and more definitive conclusions. Additionally, we are exploring how our findings could influence other training methods for MDPs, particularly in the realm of model-based reinforcement learning with continuous action and symmetry groups. We recognize that our specific focus might limit direct applicability to other training methods; however, we see potential in areas such as imitation learning when demonstrations are available.
    
    Additionally, we restate our contributions from the empirical study:
    
    1. **Demonstration of Euclidean Symmetry in Model-Based RL Algorithms:**
        - The paper defines a subclass of Markov Decision Processes (MDPs), termed Geometric MDPs, that are prevalent in robotics and exhibit additional structural properties. These MDPs, when linearized, adhere to steerable kernel constraints, leading to a substantial reduction in the parameter space.
    2. **Development of a Novel Model-Based RL Algorithm:**
        - Utilizing the insights gained from the theoretical framework, the paper develops a model-based RL algorithm that leverages Euclidean symmetry. This algorithm is shown to outperform standard techniques in common RL benchmarks.
    3. **Consideration of Equivariance in Sampling-Based RL Methods:**
        - The paper claims to be the first to consider the importance of equivariance in sampling-based RL methods. This contribution provides a deeper understanding of symmetry in RL algorithms and offers insights for future research.
    4. **Practical Limitations and Scope:**
        - While Euclidean symmetry can lead to significant savings in parameters, it does not always offer practical benefits for tasks with local coordinates. For example, locomotion tasks do not greatly benefit from this approach.
        - The algorithm assumes that the symmetry group is known, which is typically determined by the robot's workspace dimension (often 2D or 3D).
    5. **Implications and Future Directions:**
        - The findings underscore the value of Euclidean symmetry in model-based RL algorithms, especially in robotics. However, the paper also acknowledges that further research is needed to explore the full potential and limitations of this approach, particularly in tasks with different characteristics or constraints.