Learning Policy Committees for Effective Personalization in MDPs with Diverse Tasks

Thank you for your detailed reviews and thoughtful questions!

Comment: It is still unclear to me how the proposed method addresses the issue. In the context-based meta-RL method, the context of diverse tasks is learned, which can target different patterns for diverse tasks. What is the advantage of the proposed methods?

Response: This is a really insightful point. First, we note that in our experiments, VariBAD, RL2, CMTA, MOORE, and CARE are context-based methods which serve as our baselines. Thus, we provide strong empirical evidence that cleverly clustering the tasks into groups of semantically similar tasks yields significant boost in performance. Conceptually, while in principle one can learn policies that depend on arbitrary context (captured by task parameters), the practical issue is that the sample complexity of such approaches would require an enormous number of diverse training tasks, which we simply do not have in practice (or in Meta-World). In a world with a fairly limited set of tasks, grouping similar tasks together and learning a small set of distinct policies for each group performs far better, and our theoretical analysis provides a principled grounding for such an approach. We added some of this discussion in the revised Model section.

Comment: In the experiment, It is unfair to use the parametric task representation for the proposed method as the baselines of meta-RL do not use the information.

Response: We wish to clarify that our multi-task baselines CMTA, MOORE, and CARE as well as the meta-RL baselines RL2 and VariBAD also depend on task-specific context representations. We view the fact that our approach can leverage any task representation as one of its key features in comparison to much prior work in which task representation is an inextricable part of the algorithm.

Comment:Can the proposed method deal with the scenario without any parametric information about the tasks?

Response: Indeed, we devote Section 3.4 to this issue, and show that in many real tasks, such as Meta-World, we can use semantic (natural language) descriptions of the tasks and LLMs to obtain a useful task representation. Consequently, our method can be applied even to many practical non-parametric multi-task RL settings.

Comment: If the tasks' parameter can be obtained, can we directly train a general policy conditional on the parameter that can deal with all tasks in the distribution?

Response: As we mention above, doing so successfully would require far more training tasks than is available in practical benchmark, such as Meta-World. Our key insight is that we can side-step this issue by leveraging natural similarity of many real-world tasks, while still accounting for diversity among tasks classes.

Comment: How the proposed method can achieve zero-shot generalization for the new tasks?

Response: Our "zero-shot" experiments map to two potential use cases. The first is analogous to LLM use cases where multiple responses to a single prompt can (and sometimes are) provided, with the constraint that the number of such simultaneous responses must be very small (e.g., $K\le 3$, as in all our experiments). The user then chooses the most helpful response. The second is when a set of new tasks arrives after training, and either human expertise (along with task descriptions, such as using natural language), or low-cost evaluation, can quickly determine which is the most appropriate policy.

We deal with the latter setting more systematically in our "few-shot learning" experiments, which indeed demonstrate that one can quickly learn which policy among $K$ to use when facing new tasks.

Comment: Do the few-shot adaptations in all baselines require the same sample number?

Response: Yes, to ensure a fair comparison, we keep the total number of samples during few-shot adaptation the same for all approaches.

Dear Reviewer UGnN, the discussion period will end soon. We hope that you have had a chance to consider our responses. Please let us know whether they have addressed your concerns, and if you have any further questions.

Best,

Authors

Thank you for the detailed reviews and thoughtful feedback!

Comment: I wonder how your theoretical results compare with transition-kernel learning RL work [1-3] under the scenario with identical transition kernels and different rewards.

Response: Thank you for your insightful question! The main differences are that this prior work (and related papers) typically makes significant stronger assumptions. A common example is that the action space is typically assumed to be finite. The papers that do not assume this require instead linearity (of both the dynamics and rewards), and/or make assumptions (such as low-rank) on the transition model. We do not assume linearity, make no assumptions about the nature of the dynamics, and allow actions to be either discrete or in a continuous vector space (none of our results depend on the size or dimension of state or action space). Thus, one of our key results in few-shot generalization depends only on the number of policies $K$ (i.e., sample complexity is linear in $K$, along with other problem-specific features), but has no dependence on state or action space dimension.

Comment: In Gradient-Based Coverage, it seems that your cluster representation can only be selected from tasks set $T$.

Response: We wish to clarify that our gradient-based coverage approach optimizes the parameter vector in the space of all feasible task parameter vectors, not merely restricted to $T$. The key advantage of this algorithm over Greedy Intersection is that it scales far better in the dimension of the parameter space.

Dear Reviewer GRXu, the discussion period will end soon. We hope that you have had a chance to consider our responses. Please let us know whether they have addressed your concerns, and if you have any further questions.

Best,

Authors

Thank you for your detailed reviews and thoughtful questions! Please check on our new revision.

Comment: No discussion on why covering should help compared to a single policy... For example, in the mentioned paragraph the authors claim a policy committee should help when faced with outlier tasks, does theory support this?

Response: We expanded this discussion in the revised version. In particular, using the "outlier" terminology was misleading. What we really mean is that when the set of tasks is highly diverse, it is unlikely that a single policy will effectively solve all of them very effectively, as fundamnetally different skills may be required for different sets of tasks. Consequently, what can often happen in practice is that different tasks confound one another's training. That the key insight that gives rise to our approach. Moreover, if tasks are represented parametrically, this representation can be used to infer the extent to which this diversity is present (in terms of how similar, or how spread out, the tasks are in the parameter space). Our theoretical results capture this by translating it into a formal notion of "task cover", and using this concept to obtain clustering algorithms that group similar tasks together (i.e., which presumably require similar sets of skills to accomplish) and learn a distinct model that is most appropriate for a given group of tasks. While in principle it is possible to learn a single policy that takes parametric representations of tasks as an input to account for such diversity, in practice this requires a prohibitive number of sample tasks, especially if task representations are high-dimensional and tasks very diverse. Our approach thus allows a very efficient use of a relatively small sample of tasks, as demonstrated in our experiments (for example the MOORE and CARE baseline use context information).

Comment: A big limitation that was not discussed throughout the paper, both in the theoretical and empirical parts is the assumption of access to the parametric space of tasks and the mapping from that space to the space of MDPs. The access to the parametric space is needed to perform the clustering itself and the access to the mapping is needed in the theory part in order to map between the covering parameters to MDPs to learn the policies, which in turn will consist of the committee. Both of these assumptions are substantial and a discussion of them and ways to alleviate them can make the paper stronger. One possible way to do so is to use a learned parametric representation as done in previous works (e.g. [2, 3]).

Response: This is a great point, and one that we deal with practically in Section 3.4, where we suggest that for many real non-parametric multi-task environments, we can extract a representation (embedding) using natural language descriptions and pre-trained LLMs. This is conceptually inline with [2,3], but simpler, and (as we show) quite effective in practice in the case of Meta-World.

Comment: One of the biggest (not discussed) advantages of the suggested “naive” Greedy Elimination Algorithm is that it doesn’t require access to the mapping function as it only uses existing tasks.

Such feature extraction function is actually important even for Greedy Elimination, since the coverage criterion that this (and all other) algorithms we present relies on is representation similarity. Our key practical idea is that once we have identified the clusters, we can execute RL on the set of tasks comprising each cluster. We now made this point more explicitly in Section 3.3.

Comment: Some of the settings in the theoretical part did not match the ones in the empirical study. The main difference (which was not discussed) is that in the empirical study the authors suggested that each policy in the committee will be trained on all tasks in the respective cluster, whereas in the theoretical analysis, the authors assumed each policy is the optimal one for a specific task. This modification is important as in some tasks (e.g. pick and place), the Lipshitz assumption in Lemma 1 will not hold and without the modification, I would expect the method to fail. Theoretical guarantees for the settings in the empirical part would strengthen the paper.

Response: This is a very insightful point. Indeed, our theoretical analysis is slightly stylized in this sense. What we show is that it nevertheless provides us with useful practical algorithms that outperform the state of the art in the complex benchmark setting of Meta-World, both in terms of multi-task RL (in-sample performance) and meta-RL (generalization and few-shot learning). In other words, our theoretical analysis provides a principled grounding for the final practical algorithms that we develop and use.

Thank you for your detailed reviews and thoughtful questions! Please check on our new revision.

Comment: no explicit description of how to do cluster in the following context (Sec 3.2).

Response: This section is devoted to presenting our algorithmic approaches for clustering. We provide the pseudocode (Algorithm 1, in the Appendix) for our first main contribution, which is the Greedy Intersection clustering algorithm. Greedy elimination simply iteratively chooses a single $\theta$ from the set of all tasks which covers the most tasks not previously covered. Finally, the gradient-based coverage method uses conventional gradient descent to minimize the objective in Equation (2).

Comment: The algorithm relies on task representations.

Response: Indeed, this is an important assumption in our work, and we devote Section 3.4 to this issue. In particular, we show that the assumption is not as strong as it first appears to providing an approach (with further details in the Appendix) for using pretrained LLMs to obtain parametric representations (embeddings) from textual descriptions of non-parametric tasks in Meta-World, which is a major benchmark for multi-task RL and meta-RL.

Comment:: The impact of hyperparameters $K$ and $\epsilon$ on performance.

We added further ablations to evaluate the impact of these for Meta-World and updated Fig. 2 as well as added new results in Appendix Section G.2.2.

An interesting observation from our experiments is that increasing $K$ beyond what already provides full coverage for a given $\epsilon$ significantly harms performance. The reason is that it is useful to have sufficient tasks for each cluster to enable effective and robust per-cluster training.

We find that increasing $\epsilon$ to cover more tasks can also improve performance (for a similar reason that increasing $K$ may not, as higher $\epsilon$ can ensure that we do not end up with clusters with too few tasks). Of course, for sufficiently high $\epsilon$, only a single cluster will emerge, so this, too induces an interesting tradeoff.