Paper: Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts (ICLR-2024)

Reviewer comment (Reproducibility):
The paper is not clear about the computation cost of the constrained optimization problem in eq(1).

Author response:
Instead of solving the expensive constrained optimization problem at each time step, we take into account the hard constraint by applying the Gram-Schmidt process on the generated representations from the experts to orthogonalize them, hence satisfying the constraint at each time step. The time complexity of the Gram-Schmidt process is T = $O(k^{2}\times d)$, where $d$ is the representation dimension and $k$ is the number of experts.

Reviewer comment (Experiments):
The paper lacks some more experimental discussion to prove conclusively that the orthonormal representations are helping over their closest competitor MOE.

Author response:
Our experiments show that MOORE consistently outperforms MOE on the challenging MTRL tasks of MiniGrid and MetaWorld MT10. Nevertheless, in the revised version of the paper, we now provide new results on the very challenging MetaWorld MT50, where MOORE again outperforms the MOE, even by a large margin. We believe that this can convince the Reviewer about the superiority of MOORE w.r.t. MOE and about the importance of diversity in the representation.

Reviewer comment (Theory):
There is no discussion in the paper regarding the computational cost of the constrained optimization in (1). How hard is it to solve and enforce this constraint?

Author response:
Instead of solving the expensive constrained optimization problem at each time step, we take into account the hard constraint by applying the Gram-Schmidt process on the generated representations from the experts to orthogonalize them, hence satisfying the constraint at each time step. The time complexity of the Gram-Schmidt process is T = $O(k^{2}\times d)$, where $d$ is the representation dimension and $k$ is the number of experts.

Reviewer comment (Evaluation):
It was vaguely referred to in section 5.1 (below figure 4) that MOORE has a faster convergence rate than other baselines. Can you elaborate on this? Do you believe that diverse representations (the key thrust of this paper) lead to a faster convergence?

Author response:
In MTRL, sharing representations is key to learning multiple tasks using the same resources. Learning representations and skills shared across tasks can help in solving difficult tasks. We believe the MOE baseline eventually learns independent experts at the end of the training, yet it requires more samples to do so. It needs to focus on learning diverse experts and avoid any representation collapse while enhancing the quality of representation for every expert. In MOORE, we force diversity through the application of the Gram-Schmidt process, hence the network will only focus on learning high-quality representations for each expert. We argue this is the reason why our approach exhibits high sample efficiency in comparison to the other baselines.

Reviewer comment (Experiments):
The ablation study is somewhat confusing. Observe that the ablation study is done on single-head architecture and you already show in MT7 (figure 1) that MOE is outperformed by MOORE. However, PC-grad and MTPPO outperform MOORE in figure 1 (single head). So it is not clear whether even PC-grad and MTPPO are learning a diverse set of representations. Can you elaborate on this?

Author response:
In Fig. 4, we do the ablation study on MT7 since we can evaluate MOORE on a large number of experts before reaching the case where we have an expert for each task ($k=7$). Besides, we select the single-head architecture due to time reasons, since the experiments are faster than the multi-head ones. As the Reviewer points out, MOORE performs slightly worse than the baseline in this setting. However, we consider this lower performance to be negligible, especially considering the superiority of MOORE in all the other settings.

Reviewer comment (Writing):
The introduction does not state whether the proposed method is used for learning in a multi-task scenario or to curate diverse tasks for learning. Moreover, providing a motivation why a common state representation across tasks is needed could be useful. Also, the empirical studies could be included with the contributions listed in the introduction.

Author response:
In this work, our goal is to learn a set of diverse representations that cover the common objects, skills, and properties shared across tasks. We want to remark that sharing knowledge is known to be key to the success of the MTRL algorithm, as discussed in multiple works in the MTRL literature [1,2,5,6].

Thank you for your suggestion. We have added the empirical results regarding the state-of-the-art performance of MOORE on MetaWorld as part of the contribution list.

Reviewer comment (Presentation):
In Section 4, in addition to the visualization of the mixture of experts in Fig. 1, visualizing the parts of the compositional policy $\pi(a|s)=(V_sw_c)^T\theta$ and their connection would be helpful. To further elaborate on the intertwinement of the components and the proposed process, showcasing an exemplary algorithm might be useful (Section 4.2).

Author response:
Thank you for your feedback. In the revised version, we have updated the visual illustration of MOORE in Figure 1 by demonstrating in detail the computation happening inside the output head.

Reviewer comment (Experiments):
Regarding the experiments: I assume MOORE to be used in conjunction with PPO for the MiniGrid results (5.1) and with SAC for the MetaWorld results (5.2). Yet this should be clearly stated in the paper. Furthermore, if not introduced in related work, used baselines should be described (e.g., MTPPO, MTSAC, PCGrad, ...). To improve the overall readability, the Figure placement could be improved.

Author response:
We appreciate your comment. We stated that fact in Appendix A.1.2 and A.2.2. However, we have also clarified that point in Section 5 in the revised version of the paper.

All baselines are well-known in the MTRL literature. We provided a brief description of each baseline in the paper. Moreover, in the revised version, we have added the implementation details about MTPPO and MTSAC in Appendix A.1.2 and A.2.2 since they are a central part of our approach and the other related baselines. 

Thank you. We have enhanced the figure placement in the revised version of the paper.

Reviewer comment (Presentation):
Limitations should be stated.

Author response:
MOORE is considered as a soft-MOE algorithm because it requires computing the whole $k$ representations from all experts and then aggregating them, compared to the sparse selection of one expert as done in the sparse-MOE algorithms. Thus, MOORE has the limitation of potentially suffering from high time complexity, resulting in a trade-off between the representation capacity and time complexity. In this work, our goal is to make use of the whole representational space spanned by all the diverse experts and consider them as a set of basis, letting the task-encoder produces the task-specific weights to aggregate those bases. As a future work, we can investigate the possibility of selecting a few orthogonal experts like in the sparse-MOE algorithms. In the revised version of the paper, we have stated this limitation in the conclusion.

Reviewer comment (Novelty):
Related work should be extended with regards to work on mixtures of experts for reinforcement learning.

Author response:
Thank you for the additional reference. Considering the focus on MTRL in our work, we focus on reviewing related works on a mixture of experts in the domain of MTRL. Nevertheless, we have added additional references about the mixture of experts and diversity in RL in the revised version of the paper.

Reviewer comment (Experiments):
The baselines used later on should be briefly introduced and compared.

Author response:
All baselines are well-known in the MTRL literature. We provided a brief description of each baseline in the paper. Moreover, in the revised version, we have added the implementation details about MTPPO and MTSAC in Appendix A.1.2 and A.2.2 since they are a central part of our approach and the other related baselines.

Reviewer comment (Novelty):
Further related work regarding diversity-based RL (e.g., B. Eysenbach, A. Gupta, J. Ibarz, and S. Levine, ‘Diversity is All You Need: Learning Skills without a Reward Function’, in International Conference on Learning Representations, 2019.) could also be helpful.

Author response:
Thank you for the additional reference. Considering the focus on MTRL in our work, we focus on reviewing related works on a mixture of experts in the domain of MTRL. Nevertheless, we have added additional references about the mixture of experts and diversity in RL in the revised version of the paper.

Reviewer comment (Theory):
From my understanding, orthogonality is induced via the Gram-Schmidt Process. Could you elaborate on the necessity of the additional hard constraint introduced in Eq. 1?

Author response:
The constrained optimization problem in Eq. 1 shows our objective mathematically. However, we realize the hard constraint by applying the Gram-Schmidt process on the generated representations from the experts to orthogonalize them, hence satisfying the constraint. We modified the main paper adding two sentences clarifying this part. Thank you for pointing this out.

Reviewer comment (Experiments):
How many experts are used for the MetaWorld evaluations? How does the number of experts impact the scalability of the proposed approach? Further elaborations on the computational overhead introduced, as well as possible limitations, might be more insightful than the presented remarks regarding interpretability.

Author response:
For MiniGrid, the number of experts $k$ (generating the $k$ representations) is $2$, $3$, and $4$ for MT3, MT5, and MT7, respectively. On the other hand, the number of experts is $4$ and $6$ for MT10-rand and MT50-rand of the MetaWorld benchmark, respectively. We have clarified that point in the revised version of the paper.

In Fig. 4, we showed the effect of increasing the number of experts regarding the performance. On average, we can clearly see that the performance is increasing much faster than the MOE baseline. However, this comes with a cost. The more experts we use, the higher the time complexity and the memory requirement. So, there is definitely a trade-off between the representation capacity and the time and memory requirements. 

The difference between MOORE and MOE is in the Gram-Schmidt stage where we orthogonalize the $k$-representations. The time complexity of the Gram-Schmidt process is $T = O(k^{2}\times d$) [3,4], where $d$ is the representation dimension and $k$ is the number of experts. MOORE and MOE can be considered as a soft-MOE because they both compute the whole $k$ representations from all the experts and then aggregate them. MOORE and MOE have the same memory requirement of storing the parameters of the mixture of experts. The memory requirement is also similar in PaCo; however, in the MetaWorld experiments, we exhibit better performance using fewer experts. This shows that diversity plays an important role in maximizing the representation capacity. In the revised version, we have provided a clear discussion about the computation and memory requirements in Appendix C.

Reviewer comment (Presentation):
p. 2: To improve readability, the contributions could be displayed as a list.

Author response:
Thank you for this comment. We have updated the contribution paragraph in the revised version by adding an index for each contribution.

Reviewer comment (Reproducibility):
Please provide an algorithm box. I am not sure why these different vectors are called experts. Is it for every step of the convex optimization, you perform the Gram-Schmidt orthonormalization like projected gradient descent? Or do you perform the normalization towards the end?

Author response:
Algorithm box

To clear any ambiguity, we provide now an algorithm box in the appendix of the revised version, regarding the adaptation of the actor (Algorithm 1) and critic (Algorithm 2) using MOORE. We consider the state encoders as experts. At each step, the generated representations $U$ from the experts (state encoders) are orthogonalized by applying the Gram-Schmidt process to satisfy the hard constraint in Eq. 1. We aggregate the orthogonal representations $V$ using the task-specific weights $w_{c}$.

Reviewer comment (Theory):
Say you obtain $U^* = [u_1, u_2, ..., u_k]$ by minimizing $J(\Theta)$. Are these k vectors linearly independent? Gram-Schmidt process works on linearly independent vectors. Moreover, the vectors that you get after the process, $V^*= [v_1, v_2, ..., v_k]$ span the same space as $U^*$ but the $w_c$ for both will be different. I think these confusions can be cleared if you provide an algorithm box.

Author response:
Are these k vectors linearly independent? Gram-Schmidt process works on linearly independent vectors.

We agree that the input vectors to the Gram-Schmidt process should be linearly independent. In this work, we cannot guarantee the independence of the vectors before the process; however, we have not encountered any instability during learning. We believe that the reason is that the Gram-Schmidt is a differentiable process and the gradients are actually contributing to learning the experts, hence favoring the independence of the input vectors.

Reviewer comment (Evaluation):
You mention that MOORE even outperforms the single task performance is a few cases. But in those cases, most baselines also outperform the single task performance.

Author response:
You mention that MOORE even outperforms the single task performance is a few cases. But in those cases, most baselines also outperform the single task performance.

Here, we highlighted that MOORE outperforms the single-task performance with a **significant margin**. The Reviewer correctly points out that also the other baselines manage to outperform the single-task performance in the same setting, although we can observe that they do it by a lower margin than ours. We agree that we should have stated this point clearer and we have done it in the revised version.

Reviewer comment (Writing):
What do you mean by frames in Definition 4.1?

Author response:
What do you mean by frames in Definition 4.1?

We use the term orthonormal $k$-frame, in the definition, to represent a set of orthonormal vectors. If the reviewer prefers, we can replace the term "frames" with "vectors".

Reviewer comment (Reproducibility):
How do you parameterized the critic? You have only mentioned the parameterization of actor.

Author response:
How do you parameterized the critic? You have only mentioned the parameterization of actor.

We parameterize the critic similar to the actor. Only in MetaWorld, instead of mapping the state only, we map the state-action by the mixture of experts, followed by the application of the Gram-Schmidt process to orthogonalize the representations. Then, the task-encoder generates the task-specific weights $w_{c}$ given the context $c$ (e.g., as task-id). The task-specific representation $v_{c}$ of the state-action can be computed using the task-specific weight $w_{c}$ and the set of representations $V_{s,a}$. The $Q$-value is then calculated either using a single-head architecture where we concatenate the task-id with the task-specific representation as an input to a single output module $f_{\theta}$, or using a multi-head architecture where the output module consists of multiple task-specific output modules where the corresponding output module is selected using the context $c$ (e.g., as task-id). In the revised version, we provide Alg. 1 and Alg. 2 of the MOORE forward function for the Actor and of the critic, respectively.

Reviewer comment (Reproducibility):
Explain MT-PPO - How is PPO adapted for multi-task RL?

Author response:
Explain MT-PPO - How is PPO adapted for multi-task RL?
 
We adapt PPO to the MTRL setting by training the actor (policy) and critic (value function) by minimizing their corresponding loss functions averaged on transitions sampled from all tasks. We have added the PPO adaptation details in Appendix A.1.2.

Reviewer comment (Writing):
For transfering experts, you write "we transfer experts learnt on MT3 to MT5 (MT3 $\rightarrow$ MT5) and on MT5 to MT7 (MT5 $\rightarrow$ MT7) ". But in the next sentance you write "First, we train on base tasks (intersection of two scenarios), and then we transfer the learnt experts on novel tasks (the difference between the two scenarios)". Both statements sort of contradict each other.

Author response:
Transfer learning study

We agree that this part needs more elaboration. We have added Appendix B.2 in the revised version of the paper to elaborate more on the experimental details of the transfer learning study. In our transfer learning study, we assess the transfer capability of our approach in utilizing the diverse representations, learned on a set of *base* tasks, for a set of *novel* but related tasks. We evaluate our approach, MOORE, against the MOE baseline on MiniGrid. We refer to the transfer learning adaptation of our approach as **Transfer-MOORE**, and **Transfer-MOE** for the MOE baseline.

We conducted two experiments based on the sets of tasks defined on MiniGrid (MT3, MT5, and MT7). We show the empirical results on two transfer learning scenarios where we transfer a set of experts learned on MT3 to MT5 (MT3 $\rightarrow$ MT5), and on MT5 to MT7 (MT5 $\rightarrow$ MT7). It is worth noting that MT3 is a subset of MT5, and MT5 is a subset of MT7. We consider the intersection between every two sets (MT3 and MT5 or MT5 and MT7) as base tasks while considering the difference as novel tasks. For instance, in the MT3$\rightarrow$MT5 scenario, the base tasks are LavaGap,  RedBlueDoors, and Memory (common for MT3 and MT5), while having DoorKey, and MultiRoom as novel tasks (only in MT5).

Transfer-MOORE learns the base tasks and then transfers the learned mixture of experts for learning the novel ones. While learning the novel tasks, the mixture of experts is frozen but the task encoder and output head are being trained. On the contrary, MOORE is only trained on novel tasks from scratch. This also holds for MOE and Transfer-MOE. In this study, we are employing a multi-head architecture, the output head consists of completely decoupled task-specific output modules. Therefore, we train only the task encoder and the task-specific output modules for the novel tasks from scratch, keeping the mixture of experts frozen and the base task-specific heads too.

Reviewer comment (Novelty):
The paper doesn't introduce a method that is fundamentally distinct from the approach in prior work [1], which utilizes a representative parameter set {$\phi_1,...,  \phi_k$} and task-specific weight $w_c$ derived from the task-id. The proposed structure for the policy appears to be merely different combinations of the existing learnable parameters discussed in paper [1].

Author response:
PaCo [1] interpolates a task-specific policy from a subspace of policies spanned by a set of parametric policies while utilizing task-specific weights $w_{c}$. Hence, PaCo searches for the task-specific policy in the parametric space, not in the feature/representation space. On the other hand, MOORE utilizes a set of experts (state-encoders for actor or state-action-encoders for critic) to generate a set of orthogonal representations that spans a representation subpsace where each task can interpolate the relevant representation. Using the task-id is not crucial in MOORE, nor in PaCo. It is just an example on how we can interpolate task-specific information. For instance, CARE [2] uses the task description instead of the task-id while doing attention mechanism to produce set of representations. Moreover, PaCo applies some tricks (loss-maskout and w-reset) to stabilize the MTRL training process, however, our approach does not employ any tricks.

Reviewer comment (Experiments):
Although the authors argue that orthogonality results in more diverse representations, there's an absence of theoretical backing or ablation studies to corroborate this claim.

Author response:
Thank you for raising this interesting point. We want to clarify that we are focusing on representing each state by a set of diverse representations, and we agree that we have not targeted the diversity at the task-level. Nevertheless, as shown in the experiemental section, the orthogonality condition succeeds to make the representations at the state-level diverse and outperforming the other baselines. In addition, in Fig. 4, we can deduce that the diversity manages to increase the state representation capacity as we increase the number of experts in comparison to not enforcing diversity.

Reviewer comment (Theory):
Given that the vector $k$-representations $U_s$ is contingent upon the state $s$, the authors cannot guarantee representational diversity at the task level, even if state-level orthogonality is ensured.

Author response:
Thank you for raising this interesting point. We want to clarify that we are focusing on representing each state by a set of diverse representations, and we agree that we have not targeted the diversity at the task-level. Nevertheless, as shown in the experiemental section, the orthogonality condition succeeds to make the representations at the state-level diverse and outperforming the other baselines. In addition, in Fig. 4, we can deduce that the diversity manages to increase the state representation capacity as we increase the number of experts in comparison to not enforcing diversity.

Reviewer comment (Experiments):
The authors should benchmark their method against the MT50 benchmark from Meta-World, as all their experiments are conducted on a limited task set size.

Author response:
We agree about the importance of testing MOORE on MT50, and we apologize for not having done it in the original submission due to lack of time. We have added the comparison between MOORE and the MOE baseline on MT50 for 10 epochs in the experimental section of the revised version. The experiments are still running to obtain results for 20 epochs, which we plan to publish in a new revision in the upcoming days, in order to be able to compare with the already available results of the other baselines in [1]. Notably, as shown in the new Figure 6b, the performance of MOORE obtained after 10 epochs is already superior than the state-of-the-art of PaCo [1], i.e., $57.3%$. Moreover, our performance has been achieved using only $6$ experts in comparison to the state-of-the-art PaCo which uses $20$ experts.